Fix JAX vectorize.

keras-team · Apr 22, 2024 · 99fdbf7 · 99fdbf7
1 parent 4919ea1
commit 99fdbf7
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/keras/src/backend/jax/numpy.py b/keras/src/backend/jax/numpy.py
@@ -1024,6 +1024,8 @@ def vstack(xs):
 
 
 def vectorize(pyfunc, *, excluded=None, signature=None):
+    if excluded is None:
+        excluded = set()
     return jnp.vectorize(pyfunc, excluded=excluded, signature=signature)