Don't pass dtype to lax_internal._zero

PiperOrigin-RevId: 719119309
jax-ml · Jan 24, 2025 · aa46a56 · aa46a56
1 parent 8e1f956
commit aa46a56
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/jax/_src/numpy/array_methods.py b/jax/_src/numpy/array_methods.py
@@ -818,7 +818,7 @@ def apply(self, func, *, indices_are_sorted=False, unique_indices=False,
     def _scatter_apply(x, indices, y, dims, **kwargs):
       return lax.scatter_apply(x, indices, func, dims, update_shape=y.shape, **kwargs)
     return scatter._scatter_update(self.array, self.index,
-                                   lax_internal._zero(self.array.dtype),
+                                   lax_internal._zero(self.array),
                                    _scatter_apply,
                                    indices_are_sorted=indices_are_sorted,
                                    unique_indices=unique_indices, mode=mode)