added a function

Sam-Armstrong · Nov 4, 2023 · 589f773 · 589f773
1 parent 981e3d0
commit 589f773
Showing 1 changed file with 11 additions and 1 deletion.
diff --git a/dummy-files/test.py b/dummy-files/test.py
@@ -82,4 +82,14 @@ def meshgrid(
     indexing: str = "xy",
     out: Optional[JaxArray] = None,
 ) -> List[JaxArray]:
-    return jnp.meshgrid(*arrays, sparse=sparse, indexing=indexing)
+    return jnp.meshgrid(*arrays, sparse=sparse, indexing=indexing)
+
+def ones_like(
+    x: JaxArray,
+    /,
+    *,
+    dtype: jnp.dtype,
+    device: jaxlib.xla_extension.Device = None,
+    out: Optional[JaxArray] = None,
+) -> JaxArray:
+    return jnp.ones_like(x, dtype=dtype)