Add a variant of safe_map() that has no return value, named foreach().

This avoids a bunch of list bookkeeping in cases where we are iterating only for a side effect and do not care about the results.

I would have named this iter() to match OCaml's list.iter(), but unfortunately iter() is a Python builtin.

PiperOrigin-RevId: 736859418

Author: hawkinsp

jax/_src/checkify.py

@@ -51,7 +51,7 @@
 from jax._src.tree_util import tree_unflatten
 from jax._src.typing import Array
 from jax._src.util import (as_hashable_function, split_list, safe_map, safe_zip,
-                           unzip3, weakref_lru_cache, HashableWrapper)
+                           unzip3, weakref_lru_cache, HashableWrapper, foreach)
 
 source_info_util.register_exclusion(__file__)
 traceback_util.register_exclusion(__file__)
@@ -413,8 +413,8 @@ def read_env(var: core.Atom):
   def write_env(var: core.Var, val: Any):
     env[var] = val
 
-  map(write_env, jaxpr.constvars, consts)
-  map(write_env, jaxpr.invars, in_args)
+  foreach(write_env, jaxpr.constvars, consts)
+  foreach(write_env, jaxpr.invars, in_args)
 
   # interpreter loop
   for eqn in jaxpr.eqns:
@@ -427,7 +427,7 @@ def write_env(var: core.Var, val: Any):
       error, outvals = checkify_rule(error, enabled_errors,
                                      *invals, **eqn.params)
     if eqn.primitive.multiple_results:
-      map(write_env, eqn.outvars, outvals)
+      foreach(write_env, eqn.outvars, outvals)
     else:
       write_env(eqn.outvars[0], outvals)
     core.clean_up_dead_vars(eqn, env, last_used)

jax/_src/core.py

@@ -50,7 +50,7 @@
 from jax._src.util import (safe_zip, safe_map, curry, tuple_insert,
                            tuple_delete, cache,
                            HashableFunction, HashableWrapper, weakref_lru_cache,
-                           partition_list, StrictABCMeta)
+                           partition_list, StrictABCMeta, foreach)
 import jax._src.pretty_printer as pp
 from jax._src.named_sharding import NamedSharding
 from jax._src.lib import jax_jit
@@ -578,8 +578,8 @@ def write(v: Var, val: Any) -> None:
     env[v] = val
 
   env: dict[Var, Any] = {}
-  map(write, jaxpr.constvars, consts)
-  map(write, jaxpr.invars, args)
+  foreach(write, jaxpr.constvars, consts)
+  foreach(write, jaxpr.invars, args)
   lu = last_used(jaxpr)
   for eqn in jaxpr.eqns:
     subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params)
@@ -589,7 +589,7 @@ def write(v: Var, val: Any) -> None:
         traceback, name_stack=name_stack), eqn.ctx.manager:
       ans = eqn.primitive.bind(*subfuns, *map(read, eqn.invars), **bind_params)
     if eqn.primitive.multiple_results:
-      map(write, eqn.outvars, ans)
+      foreach(write, eqn.outvars, ans)
     else:
       write(eqn.outvars[0], ans)
     clean_up_dead_vars(eqn, env, lu)
@@ -2837,7 +2837,7 @@ def write(v: Var, a: AbstractValue) -> None:
 
       # Check out_type matches the let-binders' annotation (after substitution).
       out_type = substitute_vars_in_output_ty(out_type, eqn.invars, eqn.outvars)
-      map(write, eqn.outvars, out_type)
+      foreach(write, eqn.outvars, out_type)
 
     except JaxprTypeError as e:
       ctx, settings = ctx_factory()
@@ -2848,7 +2848,7 @@ def write(v: Var, a: AbstractValue) -> None:
       raise JaxprTypeError(msg, eqn_idx) from None
 
   # TODO(mattjj): include output type annotation on jaxpr and check it here
-  map(read, jaxpr.outvars)
+  foreach(read, jaxpr.outvars)
 
 def check_type(
     ctx_factory: Callable[[], tuple[JaxprPpContext, JaxprPpSettings]],

jax/_src/interpreters/ad.py

@@ -39,7 +39,7 @@
 from jax._src.dtypes import dtype, float0
 from jax._src.util import (unzip2, safe_map, safe_zip, split_list, wrap_name,
                            as_hashable_function, weakref_lru_cache,
-                           partition_list, subs_list2)
+                           partition_list, subs_list2, foreach)
 
 zip = safe_zip
 map = safe_map
@@ -344,10 +344,10 @@ def write_primal(v, val):
       primal_env[v] = val
 
   primal_env: dict[Any, Any] = {}
-  map(write_primal, jaxpr.constvars, consts)
+  foreach(write_primal, jaxpr.constvars, consts)
   # FIXME: invars can contain both primal and tangent values, and this line
   #        forces primal_in to contain UndefinedPrimals for tangent values!
-  map(write_primal, jaxpr.invars, primals_in)
+  foreach(write_primal, jaxpr.invars, primals_in)
 
   # Start with a forward pass to evaluate any side-effect-free JaxprEqns that
   # only operate on primals. This is required to support primitives with
@@ -367,15 +367,15 @@ def write_primal(v, val):
         traceback, name_stack=name_stack), eqn.ctx.manager:
       ans = eqn.primitive.bind(*subfuns, *map(read_primal, eqn.invars), **bind_params)
     if eqn.primitive.multiple_results:
-      map(write_primal, eqn.outvars, ans)
+      foreach(write_primal, eqn.outvars, ans)
     else:
       write_primal(eqn.outvars[0], ans)
 
   ct_env: dict[Any, Any] = {}
   ctx = (source_info_util.transform_name_stack('transpose') if transform_stack
          else contextlib.nullcontext())
   with ctx:
-    map(partial(write_cotangent, 'outvars'), jaxpr.outvars, cotangents_in)
+    foreach(partial(write_cotangent, 'outvars'), jaxpr.outvars, cotangents_in)
     for eqn in lin_eqns[::-1]:
       if eqn.primitive.ref_primitive:
         if eqn.primitive is core.mutable_array_p:
@@ -417,7 +417,7 @@ def write_primal(v, val):
             raise e from None
         cts_out = [Zero(v.aval) for v in eqn.invars] if cts_out is Zero else cts_out
         # FIXME: Some invars correspond to primals!
-        map(partial(write_cotangent, eqn.primitive), eqn.invars, cts_out)
+        foreach(partial(write_cotangent, eqn.primitive), eqn.invars, cts_out)
 
   cotangents_out = map(read_cotangent, jaxpr.invars)
   return cotangents_out

jax/_src/interpreters/mlir.py

@@ -53,6 +53,7 @@
 from jax._src.sharding_impls import (AUTO, NamedSharding,
                                      modify_sdy_sharding_wrt_axis_types,
                                      SdyArraySharding, SdyArrayShardingList)
+from jax._src.util import foreach
 from jax._src.lib import xla_client as xc
 from jax._src.lib import xla_extension, xla_extension_version
 from jax._src.lib.mlir import dialects, ir, passmanager
@@ -1941,8 +1942,8 @@ def get_override_lowering_rule(primitive: core.Primitive) -> LoweringRule | None
   assert len(args) == len(jaxpr.invars), (jaxpr, args)
   assert len(consts) == len(jaxpr.constvars), (jaxpr, consts)
   assert len(ctx.shape_poly_state.dim_vars) == len(dim_var_values), (ctx.shape_poly_state.dim_vars, dim_var_values)
-  map(write, jaxpr.constvars, consts)
-  map(write, jaxpr.invars, args)
+  foreach(write, jaxpr.constvars, consts)
+  foreach(write, jaxpr.invars, args)
   last_used = core.last_used(jaxpr)
   for eqn in jaxpr.eqns:
     in_nodes = map(read, eqn.invars)
@@ -2009,7 +2010,7 @@ def get_override_lowering_rule(primitive: core.Primitive) -> LoweringRule | None
                        f"{eqn}, got output {ans}") from e
 
     assert len(ans) == len(eqn.outvars), (ans, eqn)
-    map(write, eqn.outvars, out_nodes)
+    foreach(write, eqn.outvars, out_nodes)
     core.clean_up_dead_vars(eqn, env, last_used)
   return tuple(read(v) for v in jaxpr.outvars), tokens
 
@@ -2101,11 +2102,11 @@ def lower_per_platform(ctx: LoweringRuleContext,
   # If there is a single rule left just apply the rule, without conditionals.
   if len(kept_rules) == 1:
     output = kept_rules[0](ctx, *rule_args, **rule_kwargs)
-    map(
+    foreach(
         lambda o: wrap_compute_type_in_place(ctx, o.owner),
         filter(_is_not_block_argument, flatten_ir_values(output)),
     )
-    map(
+    foreach(
         lambda o: wrap_xla_metadata_in_place(ctx, o.owner),
         flatten_ir_values(output),
     )
@@ -2146,11 +2147,11 @@ def lower_per_platform(ctx: LoweringRuleContext,
       except TypeError as e:
         raise ValueError("Output of translation rule must be iterable: "
                         f"{description}, got output {output}") from e
-      map(
+      foreach(
           lambda o: wrap_compute_type_in_place(ctx, o.owner),
           filter(_is_not_block_argument, out_nodes),
       )
-      map(
+      foreach(
           lambda o: wrap_xla_metadata_in_place(ctx, o.owner),
           out_nodes,
       )

jax/_src/interpreters/partial_eval.py

@@ -47,7 +47,7 @@
 from jax._src.util import (unzip2, safe_zip, safe_map, toposort, split_list,
                            merge_lists, partition_list, OrderedSet,
                            as_hashable_function, weakref_lru_cache, subs_list,
-                           HashableFunction)
+                           HashableFunction, foreach)
 
 
 map, unsafe_map = safe_map, map
@@ -1085,8 +1085,8 @@ def has_effects(effects) -> bool:
 
   newvar = core.gensym(suffix='_offload')
   known_eqns, staged_eqns = [], []
-  map(write, in_unknowns, in_inst, jaxpr.invars)
-  map(partial(write, False, True), jaxpr.constvars)
+  foreach(write, in_unknowns, in_inst, jaxpr.invars)
+  foreach(partial(write, False, True), jaxpr.constvars)
   for eqn in jaxpr.eqns:
     unks_in, inst_in = unzip2(map(read, eqn.invars))
     rule = partial_eval_jaxpr_custom_rules.get(eqn.primitive)
@@ -1098,18 +1098,18 @@ def has_effects(effects) -> bool:
           residual_refs.add(r)
         else:
           residuals.add(r)
-      map(write, unks_out, inst_out, eqn.outvars)
+      foreach(write, unks_out, inst_out, eqn.outvars)
     elif any(unks_in):
       inputs = map(ensure_instantiated, inst_in, eqn.invars)
       staged_eqns.append(eqn.replace(invars=inputs))
-      map(partial(write, True, True), eqn.outvars)
+      foreach(partial(write, True, True), eqn.outvars)
     else:
       known_eqns.append(eqn)
       # If it's an effectful primitive, we always to run and avoid staging it.
       policy = ensure_enum(saveable(
           eqn.primitive, *[x.aval for x in eqn.invars], **eqn.params))
       if has_effects(eqn.effects) or isinstance(policy, SaveableType):
-        map(partial(write, False, False), eqn.outvars)
+        foreach(partial(write, False, False), eqn.outvars)
       elif isinstance(policy, Offloadable):
         # TODO(slebedev): This is a legit error which requires a BUILD fix.
         from jax._src.dispatch import device_put_p, TransferToMemoryKind, CopySemantics  # pytype: disable=import-error
@@ -1124,7 +1124,7 @@ def has_effects(effects) -> bool:
             JaxprEqnContext(None, False))
         known_eqns.append(offload_eqn)
         # resvars are known and available in the backward jaxpr.
-        map(partial(write, False, True), resvars)
+        foreach(partial(write, False, True), resvars)
         residuals.update(resvars)
         reload_eqn = core.JaxprEqn(
             resvars, eqn.outvars, device_put_p,
@@ -1135,12 +1135,12 @@ def has_effects(effects) -> bool:
             JaxprEqnContext(None, False))
         staged_eqns.append(reload_eqn)
         # outvars are known and available in the backward jaxpr.
-        map(partial(write, False, True), eqn.outvars)
+        foreach(partial(write, False, True), eqn.outvars)
       else:
         assert isinstance(policy, RecomputeType)
         inputs = map(ensure_instantiated, inst_in, eqn.invars)
         staged_eqns.append(eqn.replace(invars=inputs))
-        map(partial(write, False, True), eqn.outvars)
+        foreach(partial(write, False, True), eqn.outvars)
   unzipped = unzip2(map(read, jaxpr.outvars))
   out_unknowns, out_inst = list(unzipped[0]), list(unzipped[1])
   assert all(type(v) is Var for v in residuals), residuals
@@ -1441,14 +1441,14 @@ def write(x: Atom, b: bool) -> None:
       env[x] = read(x) or b
 
   new_eqns = []
-  map(write, jaxpr.outvars, used_outputs)
+  foreach(write, jaxpr.outvars, used_outputs)
   for eqn in jaxpr.eqns[::-1]:
     used_outs = map(read, eqn.outvars)
     rule = dce_rules.get(eqn.primitive, _default_dce_rule)
     used_ins, new_eqn = rule(used_outs, eqn)
     if new_eqn is not None:
       new_eqns.append(new_eqn)
-    map(write, eqn.invars, used_ins)
+    foreach(write, eqn.invars, used_ins)
   used_inputs = map(read, jaxpr.invars)
   used_inputs = map(op.or_, instantiate, used_inputs)
 
@@ -2495,15 +2495,15 @@ def read(x):
   def write(v, val) -> None:
     env[v] = val
 
-  map(write, jaxpr.constvars, consts)
-  map(write, jaxpr.invars, args)
+  foreach(write, jaxpr.constvars, consts)
+  foreach(write, jaxpr.invars, args)
   last_used = core.last_used(jaxpr)
   for eqn in jaxpr.eqns:
     in_avals  = [_substitute_axis_sizes(env, v.aval) for v in eqn.invars]
     out_avals = [_substitute_axis_sizes(env, v.aval) for v in eqn.outvars]
     rule = padding_rules[eqn.primitive]
     outs = rule(in_avals, out_avals, *map(read, eqn.invars), **eqn.params)
-    map(write, eqn.outvars, outs)
+    foreach(write, eqn.outvars, outs)
     core.clean_up_dead_vars(eqn, env, last_used)
   return map(read, jaxpr.outvars)
 
@@ -2580,7 +2580,7 @@ def inline_jaxpr_into_trace(
     src_ = (src if not eqn.source_info.name_stack else
             src.replace(name_stack=src.name_stack + eqn.source_info.name_stack))
     trace.frame.add_eqn(eqn.replace(invars, outvars, source_info=src_))
-    map(env.setdefault, eqn.outvars, outvars)
+    foreach(env.setdefault, eqn.outvars, outvars)
 
   tracer_env: dict[Var, Any] = dict(zip([*jaxpr.constvars, *jaxpr.invars],
                                         [*consts, *arg_tracers]))

jax/_src/lax/lax.py

@@ -71,7 +71,8 @@
                                      PartitionSpec as P, canonicalize_sharding)
 from jax._src.typing import Array, ArrayLike, DimSize, DuckTypedArray, DTypeLike, Shape
 from jax._src.util import (NumpyComplexWarning, cache, canonicalize_axis,
-                           safe_map, safe_zip, split_list, weakref_lru_cache)
+                           safe_map, safe_zip, split_list, weakref_lru_cache,
+                           foreach)
 
 _max = builtins.max
 _min = builtins.min
@@ -106,7 +107,7 @@ def _check_static_shape(shape: Shape):
     # pass dynamic shapes through unchecked
     return
   else:
-    map(_check_static_shape, shapes)
+    foreach(_check_static_shape, shapes)
 
 def _try_broadcast_shapes(*shapes: tuple[int, ...], name: str) -> tuple[int, ...]:
   """

jax/_src/pallas/fuser/fusable_dtype.py

@@ -37,6 +37,7 @@
 from jax._src.pallas.fuser.fusable import fusable_p
 from jax._src.state import discharge as state_discharge
 from jax._src.state import primitives as state_primitives
+from jax._src.util import foreach
 
 # TODO(sharadmv): Enable type checking.
 # mypy: ignore-errors
@@ -216,11 +217,11 @@ def read_env(var: core.Atom):
   def write_env(var: core.Var, val: Any):
     env[var] = val
 
-  map(write_env, jaxpr.constvars, consts)
+  foreach(write_env, jaxpr.constvars, consts)
   assert len(jaxpr.invars) == len(
       args
   ), f"Length mismatch: {jaxpr.invars} != {args}"
-  map(write_env, jaxpr.invars, args)
+  foreach(write_env, jaxpr.invars, args)
 
   for eqn in jaxpr.eqns:
     invals = list(map(read_env, eqn.invars))
@@ -248,7 +249,7 @@ def write_env(var: core.Var, val: Any):
 
     if eqn.primitive.multiple_results:
       assert len(outvals) == len(eqn.outvars)
-      map(write_env, eqn.outvars, outvals)
+      foreach(write_env, eqn.outvars, outvals)
     else:
       write_env(eqn.outvars[0], outvals)
 

jax/_src/pallas/hlo_interpreter.py

@@ -40,6 +40,7 @@
 from jax._src.state import discharge as state_discharge
 from jax._src import util
 from jax._src.util import (
+    foreach,
     safe_map,
     safe_zip,
     split_list,
@@ -199,8 +200,8 @@ def write(v: jax_core.Var, val: Any) -> None:
     env[v] = val
 
   env: dict[jax_core.Var, Any] = {}
-  map(write, jaxpr.constvars, consts)
-  map(write, jaxpr.invars, args)
+  foreach(write, jaxpr.constvars, consts)
+  foreach(write, jaxpr.invars, args)
   lu = jax_core.last_used(jaxpr)
   for eqn in jaxpr.eqns:
     in_vals = map(read, eqn.invars)
@@ -216,7 +217,7 @@ def write(v: jax_core.Var, val: Any) -> None:
         subfuns, bind_params = eqn.primitive.get_bind_params(eqn.params)
         ans = eqn.primitive.bind(*subfuns, *in_vals, **bind_params)
     if eqn.primitive.multiple_results:
-      map(write, eqn.outvars, ans)
+      foreach(write, eqn.outvars, ans)
     else:
       write(eqn.outvars[0], ans)
     jax_core.clean_up_dead_vars(eqn, env, lu)

jax/_src/pallas/mosaic/lowering.py

@@ -65,6 +65,7 @@
 from jax._src.state.types import RefBitcaster, RefReshaper
 from jax._src.state.utils import dtype_bitwidth
 from jax._src.typing import Array, DTypeLike
+from jax._src.util import foreach
 from jax._src.util import safe_map
 from jax._src.util import safe_zip
 from jax._src.util import split_list
@@ -950,7 +951,7 @@ def write_env(var: jax_core.Var, val):
 
   for invar, bs in zip(jaxpr.invars, ctx.block_shapes):
     block_shape_env[invar] = bs
-  map(write_env, jaxpr.invars, args)
+  foreach(write_env, jaxpr.invars, args)
 
   initial_name_stack = [scope.name for scope in ctx.name_stack.stack]
   current_name_stack: list[str] = []
@@ -1011,7 +1012,7 @@ def write_env(var: jax_core.Var, val):
             f"{eqn.primitive.name}. "
             "Please file an issue on https://github.com/jax-ml/jax/issues.")
       if eqn.primitive.multiple_results:
-        map(write_env, eqn.outvars, ans)
+        foreach(write_env, eqn.outvars, ans)
       else:
         write_env(eqn.outvars[0], ans)