fix units arg to replace (#2)

Signed-off-by: nstarman <[email protected]>
GalacticDynamics · Dec 31, 2023 · c1aeaec · c1aeaec
1 parent 6660018
commit c1aeaec
Show file tree

Hide file tree

Showing 2 changed files with 80 additions and 80 deletions.
diff --git a/src/jax_quantity/_register_dispatches.py b/src/jax_quantity/_register_dispatches.py
@@ -17,7 +17,7 @@ def dispatcher(f: T) -> T:  # TODO: figure out mypy stub issue.
 
 @dispatcher
 def empty_like(x: Quantity, /, *, dtype: Any = None, device: Any = None) -> Quantity:
-    out = Quantity(jnp.empty_like(x.value, dtype=dtype), units=x.unit)
+    out = Quantity(jnp.empty_like(x.value, dtype=dtype), unit=x.unit)
     return jax.device_put(out, device=device)
 
 
@@ -30,17 +30,17 @@ def full_like(
     dtype: Any = None,
     device: Any = None,
 ) -> Quantity:
-    out = Quantity(jnp.full_like(x.value, fill_value, dtype=dtype), units=x.unit)
+    out = Quantity(jnp.full_like(x.value, fill_value, dtype=dtype), unit=x.unit)
     return jax.device_put(out, device=device)
 
 
 @dispatcher
 def ones_like(x: Quantity, /, *, dtype: Any = None, device: Any = None) -> Quantity:
-    out = Quantity(jnp.ones_like(x.value, dtype=dtype), units=x.unit)
+    out = Quantity(jnp.ones_like(x.value, dtype=dtype), unit=x.unit)
     return jax.device_put(out, device=device)
 
 
 @dispatcher
 def zeros_like(x: Quantity, /, *, dtype: Any = None, device: Any = None) -> Quantity:
-    out = Quantity(jnp.zeros_like(x.value, dtype=dtype), units=x.unit)
+    out = Quantity(jnp.zeros_like(x.value, dtype=dtype), unit=x.unit)
     return jax.device_put(out, device=device)