Add standard pbroadcast rules to more primitives. This should cover all primitives from which shard_map registered standard_rewrite rules

PiperOrigin-RevId: 741516445

Author: yashk2810

jax/_src/core.py

@@ -1528,7 +1528,7 @@ def check_valid_jaxtype(x):
 
 def update_aval_with_sharding(aval, sharding):
   if isinstance(sharding, NamedSharding):
-    aval = aval.update(sharding=NamedSharding(
+    return aval.update(sharding=NamedSharding(
         sharding.mesh.abstract_mesh,
         sharding.spec._normalized_spec_for_aval(aval.ndim)))
   return aval
@@ -1659,8 +1659,10 @@ def physical_aval(aval):
     elt_aval = physical_element_aval(aval.dtype)
     if isinstance(aval, ShapedArray):
       from jax._src.sharding_impls import physical_sharding  # type: ignore
+      vma = aval.vma if config.varying_axes_in_types.value else frozenset()
       return ShapedArray((*aval.shape, *elt_aval.shape), elt_aval.dtype,
-                         sharding=physical_sharding(aval, aval.sharding))
+                         sharding=physical_sharding(aval, aval.sharding),
+                         vma=vma)
     return DShapedArray((*aval.shape, *elt_aval.shape), elt_aval.dtype)
   return aval
 
@@ -2019,6 +2021,8 @@ def standard_insert_pbroadcast(*args):
           if out_vma - src else arg for arg, src in zip(args, in_vma)]
 
 def standard_vma_rule(prim_name, *avals, **kwargs):
+  if not avals:
+    return avals
   vma, *vmas = [a.vma for a in avals]
   if not all(vma == vma_ for vma_ in vmas):
     raise ValueError(

jax/_src/ffi.py

@@ -24,6 +24,7 @@
 
 import jax
 from jax._src import core
+from jax._src import config
 from jax._src import deprecations
 from jax._src import dispatch
 from jax._src import effects
@@ -515,7 +516,7 @@ def wrapped(*args: ArrayLike, **kwargs: Any):
               "and an output with a different layout "
               f"{static_output_layouts[o_idx]}.")
         static_input_output_aliases += ((i_idx, o_idx),)
-
+    args = core.standard_insert_pbroadcast(*args)
     results = ffi_call_p.bind(
         *args,
         result_avals=result_avals,
@@ -638,9 +639,11 @@ def ffi_call_abstract_eval(
     has_side_effect: bool,
     **_,
 ):
-  del avals_in  # unused
+  out_vma = (core.standard_vma_rule('ffi_call', *avals_in)
+             if config.varying_axes_in_types.value else frozenset())
   effects = {_FfiEffect} if has_side_effect else core.no_effects
-  return result_avals, effects
+  return tuple(r if r is core.abstract_token else r.update(vma=out_vma)
+               for r in result_avals), effects
 
 
 def ffi_call_jvp(*args, target_name, **_):

jax/_src/lax/control_flow/solves.py

@@ -23,6 +23,7 @@
 from jax._src import api
 from jax._src import api_util
 from jax._src import core
+from jax._src import config
 from jax._src import custom_derivatives
 from jax._src import linear_util as lu
 from jax._src.interpreters import ad
@@ -309,24 +310,25 @@ def f_aux(x):
   jaxprs = _LinearSolveTuple(
       matvec_jaxpr, vecmat_jaxpr, solve_jaxpr, tr_solve_jaxpr)
 
-  out_flat = linear_solve_p.bind(
-      *(_flatten(all_consts) + b_flat),
-      const_lengths=const_lengths, jaxprs=jaxprs)
+  args = _flatten(all_consts) + b_flat
+  args = core.standard_insert_pbroadcast(*args)
+  out_flat = linear_solve_p.bind(*args, const_lengths=const_lengths, jaxprs=jaxprs)
 
   return tree_unflatten(out_tree, out_flat)
 
 
 def _linear_solve_abstract_eval(*args, const_lengths, jaxprs):
   args_to_raise = args[sum(const_lengths):]
-
   # raise aux_args to shaped arrays as well if present
   # number of aux args is the difference in out_avals
   # of solve and matvec (since they map to the same vector space)
-
   num_aux = len(jaxprs.solve.out_avals) - len(jaxprs.matvec.out_avals)
   if num_aux > 0:
     args_to_raise += tuple(jaxprs.solve.out_avals[-num_aux:])
-  return args_to_raise, jaxprs.solve.effects
+  out_vma = (core.standard_vma_rule('linear_solve', *args_to_raise)
+             if config.varying_axes_in_types.value else frozenset())
+  return (tuple(a.update(vma=out_vma) for a in args_to_raise),
+          jaxprs.solve.effects)
 
 
 def _custom_linear_solve_impl(*args, const_lengths, jaxprs):

jax/_src/lax/windowed_reductions.py

@@ -338,7 +338,7 @@ def _reduce_window_abstract_eval_rule(
   out_sharding = reduce_window_sharding_rule(
       operand_avals[0], window_dimensions, window_strides, padding,
       base_dilation, window_dilation)
-  out_vma = (core.standard_vma_rule('reduce_window', operand_avals)
+  out_vma = (core.standard_vma_rule('reduce_window', *operand_avals)
              if config.varying_axes_in_types.value else frozenset())
   return tuple(ShapedArray(out_shape, op.dtype, sharding=out_sharding,
                            vma=out_vma)

jax/_src/prng.py

@@ -178,7 +178,9 @@ def copy_to_host_async(self):
   def aval(self):
     logical_sharding = (self.sharding if hasattr(self._base_array, 'sharding')
                         else None)
-    return keys_shaped_array(self._impl, self.shape, logical_sharding)
+    vma = (self._base_array.aval.vma if config.varying_axes_in_types.value else frozenset()
+           if hasattr(self._base_array, 'aval') else frozenset())
+    return keys_shaped_array(self._impl, self.shape, logical_sharding, vma)
 
   @property
   def shape(self):
@@ -329,8 +331,8 @@ def seed_with_impl(impl: PRNGImpl, seed: int | typing.ArrayLike) -> PRNGKeyArray
   return random_seed(seed, impl=impl)
 
 
-def keys_shaped_array(impl, shape, sharding):
-  aval = core.ShapedArray(shape, KeyTy(impl))
+def keys_shaped_array(impl, shape, sharding, vma):
+  aval = core.ShapedArray(shape, KeyTy(impl), vma=vma)
   return core.update_aval_with_sharding(aval, sharding)
 
 def base_arr_shape_to_keys_shape(impl, base_arr_shape):
@@ -550,7 +552,8 @@ def random_seed(seeds: int | typing.ArrayLike, impl: PRNGImpl) -> PRNGKeyArray:
 
 @random_seed_p.def_abstract_eval
 def random_seed_abstract_eval(seeds_aval, *, impl):
-  return keys_shaped_array(impl, seeds_aval.shape, seeds_aval.sharding)
+  out_vma = seeds_aval.vma if config.varying_axes_in_types.value else frozenset()
+  return keys_shaped_array(impl, seeds_aval.shape, seeds_aval.sharding, out_vma)
 
 @random_seed_p.def_impl
 def random_seed_impl(seeds, *, impl):
@@ -584,8 +587,9 @@ def random_split_abstract_eval(keys_aval, *, shape):
   # TODO(yashkatariya): random_split should take sharding as an arg too so we
   # don't choose None here?
   new_spec = (*keys_aval.sharding.spec, *[None] * len(shape))
+  out_vma = keys_aval.vma if config.varying_axes_in_types.value else frozenset()
   return keys_shaped_array(keys_aval.dtype._impl, (*keys_aval.shape, *shape),
-                           keys_aval.sharding.with_spec(new_spec))
+                           keys_aval.sharding.with_spec(new_spec), out_vma)
 
 @random_split_p.def_impl
 def random_split_impl(keys, *, shape):
@@ -611,7 +615,9 @@ def random_split_lowering(ctx, keys, *, shape):
 
 
 def random_fold_in(keys, msgs):
-  return random_fold_in_p.bind(keys, jnp.asarray(msgs))
+  msgs = jnp.asarray(msgs)
+  keys, msgs = core.standard_insert_pbroadcast(keys, msgs)
+  return random_fold_in_p.bind(keys, msgs)
 
 random_fold_in_p = core.Primitive('random_fold_in')
 ad.defjvp_zero(random_fold_in_p)
@@ -623,7 +629,9 @@ def random_fold_in_abstract_eval(keys_aval, msgs_aval):
       'random_fold_in', keys_aval, msgs_aval)
   sharding = lax_internal.broadcasting_sharding_rule(
       'random_fold_in', keys_aval, msgs_aval)
-  return core.ShapedArray(shape, keys_aval.dtype, sharding=sharding)
+  vma = (core.standard_vma_rule('random_fold_in', keys_aval, msgs_aval)
+         if config.varying_axes_in_types.value else frozenset())
+  return core.ShapedArray(shape, keys_aval.dtype, sharding=sharding, vma=vma)
 
 @random_fold_in_p.def_impl
 def random_fold_in_impl(keys, msgs):
@@ -661,7 +669,8 @@ def random_bits(keys, bit_width, shape):
 def random_bits_abstract_eval(keys_aval, *, bit_width, shape):
   out_shape = (*keys_aval.shape, *shape)
   out_dtype = dtypes.dtype(f'uint{bit_width}')
-  return core.ShapedArray(out_shape, out_dtype)
+  vma = keys_aval.vma if config.varying_axes_in_types.value else frozenset()
+  return core.ShapedArray(out_shape, out_dtype, vma=vma)
 
 @random_bits_p.def_impl
 def random_bits_impl(keys, *, bit_width, shape):
@@ -718,7 +727,9 @@ def random_wrap(base_arr, *, impl):
 def random_wrap_abstract_eval(base_arr_aval, *, impl):
   shape = base_arr_shape_to_keys_shape(impl, base_arr_aval.shape)
   sharding = logical_sharding(shape, KeyTy(impl), base_arr_aval.sharding)
-  return keys_shaped_array(impl, shape, sharding)
+  out_vma = (base_arr_aval.vma if config.varying_axes_in_types.value else
+             frozenset())
+  return keys_shaped_array(impl, shape, sharding, out_vma)
 
 @random_wrap_p.def_impl
 def random_wrap_impl(base_arr, *, impl):