From 565bfa662dc4d91498b47d1ec9d715499670bc5c Mon Sep 17 00:00:00 2001 From: chaoming Date: Wed, 13 Nov 2024 00:21:40 +0800 Subject: [PATCH] fix tests --- brainstate/event/_xla_custom_op_test.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/brainstate/event/_xla_custom_op_test.py b/brainstate/event/_xla_custom_op_test.py index f773f4c..115d699 100644 --- a/brainstate/event/_xla_custom_op_test.py +++ b/brainstate/event/_xla_custom_op_test.py @@ -31,6 +31,8 @@ def cpu_kernel(**kwargs): def add_kernel_numba(x, y, out): out[...] = x + y + return add_kernel_numba + def gpu_kernel(x_info): return pl.pallas_call( add_vectors_kernel,