nebius/SWE-rebench-V2-PRs huggingface dataset:
Repository: google/jax
Pull request #23108 (2024-08-19 07:43:20)

Fix bug in ffi_lowering where custom layouts were ignored.
Fix bug in ffi_lowering where custom layouts were ignored.



Patch:
diff --git a/jax/_src/extend/ffi.py b/jax/_src/extend/ffi.py
index c39e05d93..3965c8b72 100644
--- a/jax/_src/extend/ffi.py
+++ b/jax/_src/extend/ffi.py
@@ -140,8 +140,12 @@ def ffi_lowering(
       kwargs["result_types"] = [mlir.aval_to_ir_type(aval) for aval in ctx.avals_out]
     if operand_layouts is None:
       kwargs["operand_layouts"] = _default_layouts(aval.shape for aval in ctx.avals_in)  # pytype: disable=attribute-error
+    else:
+      kwargs["operand_layouts"] = operand_layouts
     if result_layouts is None:
       kwargs["result_layouts"] = _default_layouts(aval.shape for aval in ctx.avals_out)
+    else:
+      kwargs["result_layouts"] = result_layouts
     if "result_shapes" not in kwargs and not all(
         core.is_constant_shape(aval.shape) for aval in ctx.avals_out):
       kwargs["result_shapes"] = [


Test patch:
diff --git a/tests/extend_test.py b/tests/extend_test.py
index 3194a3ef9..cdf3af8fb 100644
--- a/tests/extend_test.py
+++ b/tests/extend_test.py
@@ -25,6 +25,7 @@ import jax.numpy as jnp
 
 from jax._src import abstract_arrays
 from jax._src import api
+from jax._src import core
 from jax._src import linear_util
 from jax._src import prng
 from jax._src import test_util as jtu
@@ -100,6 +101,29 @@ class FfiTest(jtu.JaxTestCase):
     for header in ["c_api.h", "api.h", "ffi.h"]:
       self.assertTrue(os.path.exists(os.path.join(base_dir, header)))
 
+  def testLoweringLayouts(self):
+    # Regression test to ensure that the lowering rule properly captures
+    # layouts.
+    def lowering_rule(ctx, x):
+      aval, = ctx.avals_in
+      ndim = len(aval.shape)
+      layout = tuple(range(ndim))
+      return jex.ffi.ffi_lowering("test_ffi", operand_layouts=[layout],
+                                  result_layouts=[layout])(ctx, x)
+    prim = core.Primitive("test_ffi")
+    prim.def_impl(lambda x: x)
+    prim.def_abstract_eval(lambda x: x)
+    mlir.register_lowering(prim, lowering_rule)
+    x = jnp.linspace(0, 1, 5)
+    lowered = jax.jit(prim.bind).lower(x)
+    module = lowered.compiler_ir("stablehlo")
+    for func in module.body.operations:
+      for block in func.body.blocks:
+        for op in block.operations:
+          if op.OPERATION_NAME == "stablehlo.custom_call":
+            self.assertIn("operand_layouts", op.attributes)
+            self.assertIn("result_layouts", op.attributes)
+
   @parameterized.parameters([
       (True, mlir.ir.BoolAttr.get),
       (1, mlir.i64_attr),
Repository: google/jax
Pull request #23222 (2024-08-23 21:24:54)

add experimental jax.no_tracing context manager
I don't _think_ we need to change `trace_context` (or any of the other contexts).


Patch:
diff --git a/jax/__init__.py b/jax/__init__.py
index 037386317..168ac9278 100644
--- a/jax/__init__.py
+++ b/jax/__init__.py
@@ -56,6 +56,7 @@ from jax._src.config import (
   debug_nans as debug_nans,
   debug_infs as debug_infs,
   log_compiles as log_compiles,
+  no_tracing as no_tracing,
   explain_cache_misses as explain_cache_misses,
   default_device as default_device,
   default_matmul_precision as default_matmul_precision,
diff --git a/jax/_src/config.py b/jax/_src/config.py
index 46b327327..b6d2358f4 100644
--- a/jax/_src/config.py
+++ b/jax/_src/config.py
@@ -1501,6 +1501,11 @@ eager_pmap = bool_state(
     upgrade=True,
     help='Enable eager-mode pmap when jax_disable_jit is activated.')
 
+no_tracing = bool_state(
+    name='jax_no_tracing',
+    default=False,
+    help='Disallow tracing for JIT compilation.')
+
 disable_vmap_shmap_error = bool_state(
     name='jax_disable_vmap_shmap_error',
     default=False,
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index d8784848f..0997f1107 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -353,6 +353,9 @@ def _cpp_pjit(fun: Callable, jit_info: PjitInfo):
 
   @api_boundary
   def cache_miss(*args, **kwargs):
+    if config.no_tracing.value:
+      raise RuntimeError(f"re-tracing function {jit_info.fun_sourceinfo} for "
+                         "`jit`, but 'no_tracing' is set")
     outs, out_flat, out_tree, args_flat, jaxpr, attrs_tracked = _python_pjit_helper(
         fun, jit_info, *args, **kwargs)
     executable = _read_most_recent_pjit_call_executable(jaxpr)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index eee492443..3d3d08c09 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1463,6 +1463,20 @@ class JitTest(jtu.BufferDonationTestCase):
     self.assertAllClose(f(np.nan), np.nan)
     self.assertAllClose(jit(f)(np.nan), np.nan)
 
+  def test_no_tracing(self):
+    @jax.jit
+    def f(x):
+      return x
+
+    x = jnp.arange(3)
+    y = jnp.arange(4)
+
+    _ = f(x)  # no crash
+
+    with self.assertRaisesRegex(RuntimeError, 'no_tracing'):
+      with jax.no_tracing():
+        _ = f(y)  # crash!
+
 
 class APITest(jtu.JaxTestCase):
Repository: google/jax
Pull request #23683 (2024-09-17 04:14:51)

Make `debug.print` work with static args. Fixes: https://github.com/google/jax/issues/23600
Make `debug.print` work with static args. Fixes: https://github.com/google/jax/issues/23600



Patch:
diff --git a/jax/_src/debugging.py b/jax/_src/debugging.py
index 3e7082ab1..337349694 100644
--- a/jax/_src/debugging.py
+++ b/jax/_src/debugging.py
@@ -46,6 +46,7 @@ from jax._src.lib.mlir import ir
 from jax._src.lib.mlir.dialects import hlo
 from jax._src.sharding import Sharding
 from jax._src.sharding_impls import NamedSharding, parse_flatten_op_sharding
+from jax._src.api_util import shaped_abstractify
 from jax._src.state import discharge as state_discharge
 
 logger = logging.getLogger(__name__)
@@ -256,12 +257,29 @@ def debug_callback(callback: Callable[..., None], *args: Any,
     raise TypeError("first argument to jax.debug.callback must be callable, "
                     f"but got an object of type {type(callback)}")
   flat_args, in_tree = tree_util.tree_flatten((args, kwargs))
-  effect = ordered_debug_effect if ordered else debug_effect
-  def _flat_callback(*flat_args):
-    args, kwargs = tree_util.tree_unflatten(in_tree, flat_args)
+  static_args, dyn_args = {}, []
+  for i, a in enumerate(flat_args):
+    try:
+      shaped_abstractify(a)
+      dyn_args.append(a)
+    except (AssertionError, TypeError):
+      static_args[i] = a
+
+  def _flat_callback(*dyn_args):
+    all_args = [None] * (len(static_args) + len(dyn_args))
+    di = iter(dyn_args)
+    for i in range(len(all_args)):
+      if i in static_args:
+        all_args[i] = static_args[i]
+      else:
+        all_args[i] = next(di)
+    assert next(di, None) is None
+    args, kwargs = tree_util.tree_unflatten(in_tree, all_args)
     callback(*args, **kwargs)
     return ()
-  debug_callback_p.bind(*flat_args, callback=_flat_callback, effect=effect)
+
+  effect = ordered_debug_effect if ordered else debug_effect
+  debug_callback_p.bind(*dyn_args, callback=_flat_callback, effect=effect)
 
 class _DebugPrintFormatChecker(string.Formatter):
 


Test patch:
diff --git a/tests/debugging_primitives_test.py b/tests/debugging_primitives_test.py
index 273c12f1b..5532fdf03 100644
--- a/tests/debugging_primitives_test.py
+++ b/tests/debugging_primitives_test.py
@@ -80,6 +80,18 @@ class DebugPrintTest(jtu.JaxTestCase):
       jax.effects_barrier()
     self.assertEqual(output(), "x: 2\n")
 
+  def test_static_args(self):
+    @jax.jit
+    def f(arr):
+      jax.debug.print("arr {array}, dtype: {dtype}, arr {array2}",
+                      array=arr, dtype=arr.dtype, array2=arr)
+    arr = jnp.array([1, 2, 3], dtype=jnp.float32)
+    with jtu.capture_stdout() as output:
+      f(arr)
+      jax.effects_barrier()
+    self.assertEqual(
+        output(), "arr [1. 2. 3.], dtype: float32, arr [1. 2. 3.]\n")
+
   def test_debug_print_works_with_named_format_strings(self):
     def f(x):
       debug_print('x: {x}', x=x)
Repository: google/jax
Pull request #23704 (2024-09-17 19:05:47)

Raise a clearer error message when `closure_convert`ed function is called with inputs with the wrong structure
Fixes https://github.com/google/jax/issues/23588


Patch:
diff --git a/jax/_src/custom_derivatives.py b/jax/_src/custom_derivatives.py
index 64a37b782..019948c36 100644
--- a/jax/_src/custom_derivatives.py
+++ b/jax/_src/custom_derivatives.py
@@ -1176,7 +1176,12 @@ def _closure_convert_for_avals(fun, in_tree, in_avals):
     args, hoisted_consts = split_list(args_hconsts, [num_args])
     consts = merge(closure_consts, hoisted_consts)
     all_args, in_tree2 = tree_flatten(tuple(args))
-    assert in_tree == in_tree2
+    if in_tree != in_tree2:
+      msg = ("The inputs to the closure produced by closure_convert must have "
+             "the same Pytree structure as the example arguments passed when "
+             f"closure_convert was called. Expected {in_tree}, but got "
+             f"{in_tree2}")
+      raise TypeError(msg)
     out_flat = core.eval_jaxpr(jaxpr, consts, *all_args)
     return tree_unflatten(out_tree, out_flat)
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
index 0390e2e4b..8b75cb624 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -8988,6 +8988,19 @@ class CustomVJPTest(jtu.JaxTestCase):
     self.assertAllClose(g_c, 42. * c, check_dtypes=False)
     self.assertAllClose(g_x, 17. * x, check_dtypes=False)
 
+  def test_closure_convert_pytree_mismatch(self):
+    # See https://github.com/google/jax/issues/23588
+    def f(x, z):
+      return z * x
+
+    x, z = 2.0, 3.0
+    _, vjp = api.vjp(f, x, z)
+    vjp_pure, vjp_aux_args = jax.closure_convert(vjp, x)
+    vjp_pure(x, *vjp_aux_args)
+    with self.assertRaisesRegex(
+        TypeError, "The inputs to the closure produced by closure_convert"):
+      vjp_pure(x, vjp_aux_args)
+
   def test_float0_cotangents_automatically_handled(self):
     @jax.custom_vjp
     def f(x, y):
Repository: google/jax
Pull request #21333 (2024-05-21 20:06:40)

[Easy][djax] Better logging for name:dim mismatch
[Easy][djax] Better logging for name:dim mismatch



Patch:
diff --git a/jax/_src/interpreters/partial_eval.py b/jax/_src/interpreters/partial_eval.py
index b5ffe9ed0..9734771d3 100644
--- a/jax/_src/interpreters/partial_eval.py
+++ b/jax/_src/interpreters/partial_eval.py
@@ -2502,7 +2502,8 @@ def _complete_specs(
   for x, spec in zip(args, partial_specs):
     for i, name in spec.items():
       d = sizes.setdefault(name, x.shape[i])
-      if d is not x.shape[i] and d != x.shape[i]: raise TypeError
+      if d is not x.shape[i] and d != x.shape[i]:
+        raise TypeError(f"Provided size {d} for {name} does not match prior associated name for {name} : {x.shape[i]}")
 
   # Introduce new names as needed for Tracers in shapes.
   named_tracers: dict[TracerId, AbstractedAxisName] = {


Test patch:
diff --git a/tests/dynamic_api_test.py b/tests/dynamic_api_test.py
index 13e9cc5bb..9d5e9f671 100644
--- a/tests/dynamic_api_test.py
+++ b/tests/dynamic_api_test.py
@@ -628,6 +628,20 @@ class DynamicShapeStagingTest(jtu.JaxTestCase):
     with self.assertRaisesRegex(TypeError, msg):
       jax.make_jaxpr(jnp.ones)(jnp.ones((2, 2)))
 
+  def test_matmul_two_arg(self):
+    def f(x, y):
+      return jnp.matmul(x, y)
+
+    jaxpr = jax.make_jaxpr(f, abstracted_axes=({0: 'a_0', 1: 'a_1'}, {0: 'a_1', 1: 'a_2'}),)(jnp.ones((8, 4)), jnp.ones((4, 8)))
+
+  def test_matmul_two_arg_size_mismatch_name_validation(self):
+    def f(x, y):
+      return jnp.matmul(x, y)
+
+    with self.assertRaisesRegex(TypeError,
+                                 re.escape("Provided size 4 for a_1 does not match prior associated name for a_1 : 8")):
+      jaxpr = jax.make_jaxpr(f, abstracted_axes=({0: 'a_0', 1: 'a_1'}, {0: 'a_1', 1: 'a_2'}),)(jnp.ones((8, 4)), jnp.ones((8, 4)))
+
 @unittest.skip("Test does not work with jax.Array")
 @jtu.with_config(jax_dynamic_shapes=True, jax_numpy_rank_promotion="allow")
 class DynamicShapeAutodiffTest(jtu.JaxTestCase):
Repository: google/jax
Pull request #21368 (2024-05-22 19:58:12)

Fix `unhashable: list` error when there was a list present inside PartitionSpec.
Fix `unhashable: list` error when there was a list present inside PartitionSpec.



Patch:
diff --git a/jax/_src/sharding_impls.py b/jax/_src/sharding_impls.py
index 0bd064868..25a02facb 100644
--- a/jax/_src/sharding_impls.py
+++ b/jax/_src/sharding_impls.py
@@ -1051,7 +1051,9 @@ class ParsedPartitionSpec:
       else:
         axis_spec = (axis_spec,)
       axis_specs.append(axis_spec)
-    return cls(entry, axis_specs)
+    new_entry = PartitionSpec(
+        *[tuple(e) if isinstance(e, (list, tuple)) else e for e in entry])
+    return cls(new_entry, axis_specs)
 
   def __hash__(self):
     return hash((self.partitions, self.sync))


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index d32f46103..208e272a1 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -3495,6 +3495,12 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(cache_info2.hits, cache_info1.hits + 1)
     self.assertEqual(cache_info2.misses, cache_info1.misses)
 
+  def test_list_in_pspec(self):
+    mesh = jtu.create_global_mesh((2,), ('x',))
+    with mesh:
+      out = with_sharding_constraint(jnp.arange(8), P(['x']))
+    self.assertEqual(out.sharding, NamedSharding(mesh, P('x')))
+
   def test_sharding_preserved_trivial(self):
     mesh = jtu.create_global_mesh((2, 1), ('x', 'y'))
     ns = NamedSharding(mesh, P('x'))
Repository: google/jax
Pull request #21532 (2024-05-31 00:09:06)

Make sure that the sharding and unconstrained_dims in with_sharding_constraint are correct when wsc is vmapped.
Make sure that the sharding and unconstrained_dims in with_sharding_constraint are correct when wsc is vmapped.

In other words, if unconstrained_dims is specified, then the sharding should also contain P.UNCONSTRAINED under vmap.



Patch:
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index df5c3b5d7..c72728545 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -2481,12 +2481,22 @@ def _sharding_constraint_batcher(insert_axis, spmd_axis_name, axis_size,
   new_parts = (axis_name,) if insert_axis else (
       None if spmd_axis_name is None else spmd_axis_name)
   unconstrained_dims = {ud + (d <= ud) for ud in unconstrained_dims}
+
   if new_parts is None:
     unconstrained_dims.add(d)
+
+  vmapped_sharding = _pjit_batcher_for_sharding(
+      sharding, d, new_parts, resource_env.physical_mesh, x.ndim)
+  if unconstrained_dims and isinstance(vmapped_sharding, NamedSharding):
+    new_spec = list(vmapped_sharding.spec) + [None] * (x.ndim - len(vmapped_sharding.spec))
+    for u in unconstrained_dims:
+      new_spec[u] = PartitionSpec.UNCONSTRAINED
+    vmapped_sharding = NamedSharding(
+        vmapped_sharding.mesh, PartitionSpec(*new_spec))
+
   y = sharding_constraint_p.bind(
       x,
-      sharding=_pjit_batcher_for_sharding(
-          sharding, d, new_parts, resource_env.physical_mesh, x.ndim),
+      sharding=vmapped_sharding,
       resource_env=resource_env,
       unconstrained_dims=unconstrained_dims)
   return y, d


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 28ea6a39f..77f14ce96 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4147,6 +4147,40 @@ class ArrayPjitTest(jtu.JaxTestCase):
       self.assertNotEqual(f(1), g(1))
       self.assertEqual(g(1), h(1))
 
+  def test_wsc_vmap_unconstrained_spmd_axis_name(self):
+    def get_wsc_eqn_sharding(jaxpr):
+      for eqn in jaxpr.eqns:
+        if str(eqn.primitive) == 'sharding_constraint':
+          return eqn.params['sharding'], eqn.params['unconstrained_dims']
+      for s in core.subjaxprs(jaxpr):
+        return get_wsc_eqn_sharding(s)
+
+    mesh = jtu.create_global_mesh((2, 1), ('x', 'y'))
+    inp = jnp.ones((10, 10))
+
+    def a_function(x):
+      return with_sharding_constraint(x, NamedSharding(mesh, P(P.UNCONSTRAINED)))
+
+    def vmap_the_function_spmd(y):
+      return jax.vmap(a_function, spmd_axis_name='x')(y)
+
+    f1 = jax.jit(vmap_the_function_spmd)
+    f1(inp)  # doesn't crash
+    jaxpr1 = jax.make_jaxpr(f1)(inp)
+    s1, u1 = get_wsc_eqn_sharding(jaxpr1)
+    self.assertEqual(s1.spec, P('x', P.UNCONSTRAINED))
+    self.assertEqual(u1, {1})
+
+    def vmap_the_function_no_spmd(y):
+      return jax.vmap(a_function)(y)
+
+    f2 = jax.jit(vmap_the_function_no_spmd)
+    f2(inp)  # doesn't crash
+    jaxpr2 = jax.make_jaxpr(f2)(inp)
+    s2, u2 = get_wsc_eqn_sharding(jaxpr2)
+    self.assertEqual(s2.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
+    self.assertEqual(u2, {0, 1})
+
 
 class TempSharding(Sharding):
Repository: google/jax
Pull request #21602 (2024-06-03 17:08:40)

`input_shardings` returned from the compiled executable should match the in_tree pre DCE. Otherwise calling it raises an error saying that the length of pre dce in_tree and post dce input_shardings don't match
`input_shardings` returned from the compiled executable should match the in_tree pre DCE. Otherwise calling it raises an error saying that the length of pre dce in_tree and post dce input_shardings don't match



Patch:
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index 52b7ce792..303c5d031 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -498,6 +498,11 @@ class Compiled(Stage):
   @property
   def input_shardings(self):  # PyTree[sharding.XLACompatibleSharding]
     shardings_flat = self._executable.input_shardings()
+    # Some input shardings got DCE'd
+    if self.in_tree.num_leaves > len(shardings_flat):
+      iter_shardings_flat = iter(shardings_flat)
+      shardings_flat = [next(iter_shardings_flat) if i in self._executable._kept_var_idx
+                        else None for i in range(self.in_tree.num_leaves)]
     return tree_util.tree_unflatten(self.in_tree, shardings_flat)  # pytype: disable=attribute-error
 
   @property


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index 77f14ce96..f290b50d1 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4181,6 +4181,16 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(s2.spec, P(P.UNCONSTRAINED, P.UNCONSTRAINED))
     self.assertEqual(u2, {0, 1})
 
+  def test_aot_sharding_dce(self):
+    inp = np.arange(8)
+
+    @jax.jit
+    def f(x, y):
+      return x
+
+    input_shardings, _ = f.lower(inp, inp).compile().input_shardings
+    self.assertLen(input_shardings, 2)
+
 
 class TempSharding(Sharding):
Repository: google/jax
Pull request #21649 (2024-06-04 22:16:50)

Return `out_info` from `Lowered` since we already return `args_info`.
Return `out_info` from `Lowered` since we already return `args_info`.



Patch:
diff --git a/jax/_src/stages.py b/jax/_src/stages.py
index 303c5d031..7a0a55913 100644
--- a/jax/_src/stages.py
+++ b/jax/_src/stages.py
@@ -44,6 +44,7 @@ from jax._src import traceback_util
 from jax._src import tree_util
 from jax._src.tree_util import tree_unflatten, keystr
 from jax._src import util
+from jax._src.sharding_impls import is_unspecified_or_auto
 from jax._src.layout import Layout
 from jax._src.interpreters import mlir
 from jax._src.lib.mlir import ir
@@ -373,6 +374,12 @@ class ArgInfo:
   aval: core.AbstractValue
   donated: bool
 
+@dataclass
+class OutInfo:
+  shape: tuple[int, ...]
+  dtype: Any
+  sharding: jax.sharding.Sharding
+
 
 class Stage:
   args_info: Any  # PyTree of ArgInfo
@@ -659,6 +666,14 @@ class Lowered(Stage):
         out_tree,
         no_kwargs=no_kwargs, fun_name=fun_name, jaxpr=jaxpr)
 
+  @property
+  def out_info(self):  # PyTree of OutInfo
+    out_avals = self._lowering.compile_args["global_out_avals"]
+    out_shardings = self._lowering.compile_args["out_shardings"]
+    return self.out_tree.unflatten(
+        [OutInfo(o.shape, o.dtype, None if is_unspecified_or_auto(s) else s)
+         for o, s in zip(out_avals, out_shardings)])
+
   def compile(
       self, compiler_options: CompilerOptions | None = None) -> Compiled:
     """Compile, returning a corresponding ``Compiled`` instance."""
diff --git a/jax/stages.py b/jax/stages.py
index 0a6e6082f..e567d9264 100644
--- a/jax/stages.py
+++ b/jax/stages.py
@@ -30,4 +30,5 @@ from jax._src.stages import (
   Lowered as Lowered,
   Wrapped as Wrapped,
   ArgInfo as ArgInfo,
+  OutInfo as OutInfo,
 )


Test patch:
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index f290b50d1..f025259e0 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -4191,6 +4191,16 @@ class ArrayPjitTest(jtu.JaxTestCase):
     input_shardings, _ = f.lower(inp, inp).compile().input_shardings
     self.assertLen(input_shardings, 2)
 
+  def test_aot_out_info(self):
+    inp = np.arange(8, dtype=np.int32)
+    out_info = jax.jit(lambda x: x).lower((inp, inp)).out_info
+    self.assertEqual(out_info[0].shape, (8,))
+    self.assertEqual(out_info[1].shape, (8,))
+    self.assertEqual(out_info[0].dtype, np.int32)
+    self.assertEqual(out_info[1].dtype, np.int32)
+    self.assertEqual(out_info[0].sharding, None)
+    self.assertEqual(out_info[1].sharding, None)
+
 
 class TempSharding(Sharding):