refactor: lax.broadcast_in_dim_p (#88)

Signed-off-by: nstarman <nstarman@users.noreply.github.com>
GalacticDynamics · Sep 23, 2024 · cbd25b0 · cbd25b0
1 parent c8f711f
commit cbd25b0
Showing 1 changed file with 2 additions and 10 deletions.
diff --git a/tests/myarray.py b/tests/myarray.py
@@ -219,16 +219,8 @@ def _bitcast_convert_type_p() -> MyArray:
 
 
 @register(lax.broadcast_in_dim_p)
-def _broadcast_in_dim_p(
-    operand: MyArray,
-    *,
-    shape: Any,
-    broadcast_dimensions: Any,
-) -> MyArray:
-    return replace(
-        operand,
-        array=lax.broadcast_in_dim(operand.array, shape, broadcast_dimensions),
-    )
+def _broadcast_in_dim_p(operand: MyArray, **kwargs: Any) -> MyArray:
+    return replace(operand, array=lax.broadcast_in_dim(operand.array, **kwargs))
 
 
 # ==============================================================================