Added sharding parameter: fix for JAX 0.4.31 (#25)

patrick-kidger · Aug 1, 2024 · ef28ceb · ef28ceb
1 parent dcbadb6
commit ef28ceb
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/quax/examples/zero/_core.py b/quax/examples/zero/_core.py
@@ -64,8 +64,9 @@ def _(value: Zero, *, broadcast_dimensions, shape) -> Zero:
 
 
 @quax.register(lax.convert_element_type_p)
-def _(value: Zero, *, new_dtype, weak_type) -> Zero:
-    del weak_type
+def _(value: Zero, *, new_dtype, weak_type, sharding=None) -> Zero:
+    # sharding was added around JAX 0.4.31, it seems.
+    del weak_type, sharding
     return Zero(value.shape, new_dtype)