Reverts 6a0d8a4

PiperOrigin-RevId: 800597204

Author: yashk2810

jax/_src/lax/control_flow/loops.py

@@ -540,7 +540,12 @@ def _empty_array(prefix, length_spec, aval):
   # empty = lax.empty((*prefix, *aval.shape), aval.dtype, out_sharding=sharding)
   # return core.pvary(empty, tuple(aval.vma))
   empty = core.pvary(lax.empty2(aval.dtype), tuple(aval.vma))
-  return lax.broadcast(empty, (*prefix, *aval.shape), out_sharding=sharding)
+  out = lax.broadcast(empty, (*prefix, *aval.shape), out_sharding=sharding)
+  # TODO(yashkatariya): Maybe make this more general by passing
+  # aval.memory_space to lax.broadcast and then remove this hack?
+  if aval.memory_space != core.typeof(out).memory_space:
+    out = api.device_put(out, aval.memory_space)
+  return out
 
 eval_jaxpr_p = core.Primitive('eval_jaxpr')
 eval_jaxpr_p.multiple_results = True

jax/_src/lax/lax.py

@@ -4967,7 +4967,7 @@ def _convert_element_type_bind_with_trace(trace, args, params):
             _convert_element_type_weak_type_rule,
             _convert_element_type_sharding_rule,
             partial(core.standard_vma_rule, convert_element_type_p.name),
-            None))
+            None, None))
 ad.defjvp2(convert_element_type_p, _convert_element_type_jvp_rule)
 ad.primitive_transposes[convert_element_type_p] = _convert_element_type_transpose_rule
 
@@ -6680,7 +6680,7 @@ def _broadcast_in_dim_abstract_eval(x, *dyn_shape, shape, broadcast_dimensions,
         sharding=sharding)
     new_vma = core.standard_vma_rule('broadcast_in_dim', x)
     return core.ShapedArray(shape, x.dtype, x.weak_type, sharding=new_sharding,
-                            vma=new_vma)
+                            vma=new_vma, memory_space=x.memory_space)
   # If any BInts in shape, or Tracers in dyn_shape, produce a DShapedArray
   # (even if x is a ShapedArray)
   # TODO(mattjj): unify DShapedArray with ShapedArray, and remove this code

jax/_src/lax/linalg.py

@@ -775,8 +775,7 @@ def linalg_primitive(result_dtype, accepted_dtypes, ranks, result_shape, name,
     prim.def_abstract_eval(
       partial(lax_utils.standard_abstract_eval, prim, shape_rule, dtype_rule,
               lax_utils._standard_weak_type_rule, sharding_rule,
-              partial(core.standard_vma_rule, name),
-              None))
+              partial(core.standard_vma_rule, name), None, None))
   if supports_batching:
     batching.primitive_batchers[prim] = partial(
         batching.expand_dims_batcher, prim)

jax/_src/lax/slicing.py

@@ -2694,6 +2694,12 @@ def _scatter_spec_computation(
   return None
 
 
+def _scatter_memory_space_rule(
+    operand, indices, updates, *, update_jaxpr, update_consts,
+    dimension_numbers, indices_are_sorted, unique_indices, mode):
+  return operand.memory_space
+
+
 def _scatter_sharding_rule(
     operand, indices, updates, *, update_jaxpr, update_consts,
     dimension_numbers, indices_are_sorted, unique_indices, mode):
@@ -2905,7 +2911,8 @@ def _scatter_batching_rule(scatter_op, axis_data, batched_args, batch_dims, *,
 scatter_add_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-add',
     weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
-    vma_rule=partial(core.standard_vma_rule, 'scatter_add'))
+    vma_rule=partial(core.standard_vma_rule, 'scatter_add'),
+    memory_space_rule=_scatter_memory_space_rule)
 ad.primitive_jvps[scatter_add_p] = partial(_scatter_addsub_jvp, scatter_add_p)
 ad.primitive_transposes[scatter_add_p] = partial(_scatter_addsub_transpose_rule, scatter_add_p)
 batching.fancy_primitive_batchers[scatter_add_p] = partial(_scatter_batching_rule, scatter_add_p)
@@ -2914,7 +2921,8 @@ def _scatter_batching_rule(scatter_op, axis_data, batched_args, batch_dims, *,
 scatter_sub_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-sub',
     weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
-    vma_rule=partial(core.standard_vma_rule, 'scatter_sub')
+    vma_rule=partial(core.standard_vma_rule, 'scatter_sub'),
+    memory_space_rule=_scatter_memory_space_rule
 )
 ad.primitive_jvps[scatter_sub_p] = partial(_scatter_addsub_jvp, scatter_sub_p)
 ad.primitive_transposes[scatter_sub_p] = partial(_scatter_addsub_transpose_rule, scatter_sub_p)
@@ -2925,7 +2933,8 @@ def _scatter_batching_rule(scatter_op, axis_data, batched_args, batch_dims, *,
 scatter_mul_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-mul',
     weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
-    vma_rule=partial(core.standard_vma_rule, 'scatter_mul'))
+    vma_rule=partial(core.standard_vma_rule, 'scatter_mul'),
+    memory_space_rule=_scatter_memory_space_rule)
 
 def _scatter_mul_jvp_rhs(g, x, i, y, *, dimension_numbers,
                          indices_are_sorted, unique_indices, mode, **kw):
@@ -3056,7 +3065,8 @@ def _scatter_extremal_jvp(scatter_op, primals, tangents, update_jaxpr,
 scatter_min_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-min',
     weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
-    vma_rule=partial(core.standard_vma_rule, 'scatter_min'))
+    vma_rule=partial(core.standard_vma_rule, 'scatter_min'),
+    memory_space_rule=_scatter_memory_space_rule)
 batching.fancy_primitive_batchers[scatter_min_p] = (
   partial(_scatter_batching_rule, scatter_min_p))
 batching.skippable_batchers[scatter_min_p] = lambda _: ()
@@ -3065,7 +3075,8 @@ def _scatter_extremal_jvp(scatter_op, primals, tangents, update_jaxpr,
 scatter_max_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter-max',
     weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
-    vma_rule=partial(core.standard_vma_rule, 'scatter_max'))
+    vma_rule=partial(core.standard_vma_rule, 'scatter_max'),
+    memory_space_rule=_scatter_memory_space_rule)
 batching.fancy_primitive_batchers[scatter_max_p] = (
   partial(_scatter_batching_rule, scatter_max_p))
 batching.skippable_batchers[scatter_max_p] = lambda _: ()
@@ -3225,7 +3236,8 @@ def _scatter_transpose_rule(t, operand, indices, updates, *,
 scatter_p = standard_primitive(
     _scatter_shape_rule, _scatter_dtype_rule, 'scatter',
     weak_type_rule=_argnum_weak_type(0), sharding_rule=_scatter_sharding_rule,
-    vma_rule=partial(core.standard_vma_rule, 'scatter'))
+    vma_rule=partial(core.standard_vma_rule, 'scatter'),
+    memory_space_rule=_scatter_memory_space_rule)
 ad.primitive_jvps[scatter_p] = _scatter_jvp
 ad.primitive_transposes[scatter_p] = _scatter_transpose_rule
 batching.fancy_primitive_batchers[scatter_p] = (

jax/_src/lax/utils.py

@@ -38,13 +38,14 @@ def _argnum_weak_type(*argnums):
 
 def standard_primitive(shape_rule, dtype_rule, name,
                        weak_type_rule=None, sharding_rule=None, vma_rule=None,
-                       unreduced_rule=None):
+                       unreduced_rule=None, memory_space_rule=None):
   weak_type_rule = weak_type_rule or _standard_weak_type_rule
   prim = core.Primitive(name)
   prim.def_impl(partial(dispatch.apply_primitive, prim))
   prim.def_abstract_eval(
       partial(standard_abstract_eval, prim, shape_rule, dtype_rule,
-              weak_type_rule, sharding_rule, vma_rule, unreduced_rule))
+              weak_type_rule, sharding_rule, vma_rule, unreduced_rule,
+              memory_space_rule))
   return prim
 
 def _get_array_abstraction_level(a): return a.array_abstraction_level
@@ -140,7 +141,7 @@ def multi_mem_space_rule(prim, num_out, *avals, **kwargs):
 
 def standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule,
                            sharding_rule, vma_rule, unreduced_rule,
-                           *avals, **kwargs):
+                           memory_space_rule, *avals, **kwargs):
   for a in avals:
     if isinstance(a, state.AbstractRef):
       raise ValueError(f'Attempting to pass a Ref {a} to a primitive: '
@@ -158,7 +159,9 @@ def standard_abstract_eval(prim, shape_rule, dtype_rule, weak_type_rule,
         prim, shape_rule, dtype_rule, sharding_rule, unreduced_rule, False,
         *avals, **kwargs)
     out_vma = vma_rule(*avals, **kwargs)
-    out_mem_space = _default_memory_space_rule(prim, *avals, **kwargs)
+    out_mem_space = (_default_memory_space_rule(prim, *avals, **kwargs)
+                     if memory_space_rule is None else
+                     memory_space_rule(*avals, **kwargs))
     out_aval = core.ShapedArray(
         out_shape, out_dtype, weak_type=weak_type, sharding=out_sharding,
         vma=out_vma, memory_space=out_mem_space)

tests/memories_test.py

@@ -1642,6 +1642,19 @@ def test_fn(x_in, y_in):
     self.assertArraysEqual(x_out, x1 * x1)
     self.assertArraysEqual(y_out, y1 + y1)
 
+  def test_indexing_on_host(self):
+    @jax.jit
+    @compute_on("device_host")
+    def fn2(x):
+      x = jax.device_put(x, jax.memory.Space.Host)
+      y = jnp.ones((2, 1, 4))
+      y = jax.device_put(y, jax.memory.Space.Host)
+      z = x.at[:, 1:2, :].set(y)
+      return z
+
+    x_host = jax.device_put(jnp.ones((2,3,4)), jax.memory.Space.Host)
+    fn2(x_host)  # doesn't crash
+
   def test_compute_on_cache_miss(self):
     @jax.jit
     def f(x):
@@ -1909,8 +1922,6 @@ def g(ys, _):
     compiled_text = compiled_f.as_text()
     if compiled_text is not None:
       self.assertIn('S(5)', compiled_text)
-      self.assertNotRegex(compiled_text, r"copy-start.*S\(5\)")
-      self.assertNotRegex(compiled_text, r"copy-done.*S\(5\)")
 
     compiled_stats = compiled_f.memory_analysis()
     if compiled_stats is not None:
@@ -1948,8 +1959,6 @@ def g(ys, _):
     compiled_text = compiled_f.as_text()
     if compiled_text is not None:
       self.assertIn('S(5)', compiled_text)
-      self.assertNotRegex(compiled_text, r"copy-start.*S\(5\)")
-      self.assertNotRegex(compiled_text, r"copy-done.*S\(5\)")
       self.assertRegex(compiled_text, r"dynamic-update-slice-start.*S\(5\)")
       self.assertRegex(compiled_text, r"dynamic-update-slice-done.*S\(5\)")
       self.assertRegex(compiled_text, r"dynamic-slice-start.*S\(5\)")