Merge pull request #20314 from jakevdp:dep-maps

PiperOrigin-RevId: 617212897

Author: jax authors

jax/experimental/jax2tf/jax2tf.py

@@ -36,7 +36,6 @@
 from jax import numpy as jnp
 from jax import tree_util
 from jax import sharding
-from jax.experimental import maps
 from jax.experimental import export
 from jax.experimental.export import _export
 from jax.experimental.export import _shape_poly
@@ -54,6 +53,8 @@
 from jax._src import linear_util as lu
 from jax._src import op_shardings
 from jax._src import sharding_impls
+from jax._src import maps
+from jax._src import mesh
 from jax._src import pjit
 from jax._src import prng
 from jax._src import random as random_internal
@@ -3503,7 +3504,7 @@ def _pjit(*args: TfVal,
           jaxpr: core.ClosedJaxpr,
           in_shardings: Sequence[sharding.XLACompatibleSharding],
           out_shardings: Sequence[sharding.XLACompatibleSharding],
-          resource_env: maps.ResourceEnv,
+          resource_env: mesh.ResourceEnv,
           donated_invars,
           name: str,
           keep_unused: bool,
@@ -3535,7 +3536,7 @@ def _pjit(*args: TfVal,
 
 def _pjit_sharding_constraint(arg: TfVal, *,
                               sharding: sharding.XLACompatibleSharding,
-                              resource_env: maps.ResourceEnv,
+                              resource_env: mesh.ResourceEnv,
                               _in_avals: Sequence[core.ShapedArray],
                               _out_aval: core.ShapedArray,
                               **kwargs) -> TfVal:

jax/experimental/jax2tf/tests/jax2tf_test.py

@@ -34,13 +34,13 @@
 from jax import sharding
 from jax._src import config
 from jax._src import core
+from jax._src.maps import xmap
 from jax._src import source_info_util
 from jax._src import test_util as jtu
 from jax._src import xla_bridge as xb
 from jax.experimental import jax2tf
 from jax.experimental import export
 from jax.experimental.jax2tf.tests import tf_test_util
-from jax.experimental.maps import xmap
 from jax.experimental.shard_map import shard_map
 from jax.experimental import pjit
 from jax.sharding import PartitionSpec as P

jax/experimental/jax2tf/tests/sharding_test.py

@@ -34,12 +34,12 @@
 from jax._src import compiler
 from jax._src import config
 from jax._src import maps
+from jax._src.maps import xmap
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax import lax
 from jax.experimental import jax2tf
 from jax.experimental import pjit
-from jax.experimental.maps import xmap
 from jax.experimental.shard_map import shard_map
 from jax.sharding import NamedSharding
 from jax.sharding import Mesh

tests/compilation_cache_test.py

@@ -33,11 +33,11 @@
 from jax._src import compiler
 from jax._src import config
 from jax._src import distributed
+from jax._src.maps import xmap
 from jax._src import monitoring
 from jax._src import test_util as jtu
 from jax._src import xla_bridge
 from jax._src.lib import xla_client
-from jax.experimental.maps import xmap
 from jax.experimental.pjit import pjit
 from jax.sharding import PartitionSpec as P
 import numpy as np

tests/debug_nans_test.py

@@ -23,7 +23,8 @@
 from jax._src import api
 from jax._src import test_util as jtu
 from jax import numpy as jnp
-from jax.experimental import pjit, maps
+from jax.experimental import pjit
+from jax._src.maps import xmap
 
 from jax import config
 config.parse_flags_with_absl()
@@ -125,7 +126,7 @@ def testPmapNoNaN(self):
   @jtu.ignore_warning(message=".*is an experimental.*")
   def testXmap(self):
 
-    f = maps.xmap(
+    f = xmap(
         lambda x: 0. / x,
         in_axes=["i"],
         out_axes=["i"],

tests/debugging_primitives_test.py

@@ -20,13 +20,13 @@
 import jax
 from jax import lax
 from jax import config
-from jax.experimental import maps
 from jax.experimental import pjit
 from jax.interpreters import pxla
 from jax._src import ad_checkpoint
 from jax._src import debugging
 from jax._src import dispatch
 from jax._src import test_util as jtu
+from jax._src.maps import xmap
 import jax.numpy as jnp
 import numpy as np
 
@@ -795,7 +795,7 @@ def foo(x):
         idx = lax.axis_index('foo')
         debug_print("{idx}: {x}", idx=idx, x=x)
         return jnp.mean(x, axis=['foo'])
-      out = maps.xmap(foo, in_axes=['foo'], out_axes=[...])(x)
+      out = xmap(foo, in_axes=['foo'], out_axes=[...])(x)
       debug_print("Out: {}", out)
       return out
     mesh = jax.sharding.Mesh(np.array(jax.devices()), ['dev'])
@@ -813,8 +813,8 @@ def foo(x):
   def test_unordered_print_with_xmap(self):
     def f(x):
       debug_print("{}", x, ordered=False)
-    f = maps.xmap(f, in_axes=['a'], out_axes=None, backend='cpu',
-                  axis_resources={'a': 'dev'})
+    f = xmap(f, in_axes=['a'], out_axes=None, backend='cpu',
+             axis_resources={'a': 'dev'})
     with jax.sharding.Mesh(np.array(jax.devices()), ['dev']):
       with jtu.capture_stdout() as output:
         f(np.arange(40))

tests/jaxpr_effects_test.py

@@ -19,7 +19,6 @@
 import jax
 import jax.numpy as jnp
 from jax import lax
-from jax.experimental import maps
 from jax.experimental import pjit
 from jax._src import ad_checkpoint
 from jax._src import dispatch
@@ -32,6 +31,7 @@
 from jax._src.interpreters import ad
 from jax._src.interpreters import mlir
 from jax._src.interpreters import partial_eval as pe
+from jax._src.maps import xmap
 import numpy as np
 
 config.parse_flags_with_absl()
@@ -275,7 +275,7 @@ def f(x):
       effect_p.bind(effect=foo_effect)
       effect_p.bind(effect=bar_effect)
       return x
-    f = maps.xmap(f, in_axes=['a'], out_axes=['a'])
+    f = xmap(f, in_axes=['a'], out_axes=['a'])
     jaxpr = jax.make_jaxpr(f)(jnp.arange(jax.local_device_count()))
     self.assertSetEqual(jaxpr.effects, {foo_effect, bar_effect})
 

tests/lax_control_flow_test.py

@@ -35,12 +35,12 @@
 from jax._src import test_util as jtu
 from jax import tree_util
 from jax._src.util import unzip2
-from jax.experimental import maps
 from jax.ad_checkpoint import checkpoint as new_checkpoint, checkpoint_policies
 import jax.numpy as jnp  # scan tests use numpy
 import jax.scipy as jsp
 from jax._src.lax import control_flow as lax_control_flow
 from jax._src.lax.control_flow import for_loop
+from jax._src.maps import xmap
 
 from jax import config
 config.parse_flags_with_absl()
@@ -2712,7 +2712,7 @@ def body(carry):
         i, x = carry
         return i + 1, x + lax.psum(y, 'b')
       return lax.while_loop(cond, body, (0, z))[1]
-    maps.xmap(f, axis_sizes=dict(a=2, b=10), out_axes=(['a']), in_axes={})(1.)
+    xmap(f, axis_sizes=dict(a=2, b=10), out_axes=(['a']), in_axes={})(1.)
 
   def test_while_loop_fixed_point_with_batched_pred_and_consts(self):
     def f(i, x):

tests/pjit_test.py

@@ -41,7 +41,6 @@
 from jax.lax import with_sharding_constraint
 from jax._src import prng
 from jax.sharding import PartitionSpec as P
-from jax.experimental.maps import xmap
 from jax.experimental import multihost_utils
 from jax.experimental.custom_partitioning import custom_partitioning
 from jax._src import array
@@ -52,6 +51,7 @@
     AUTO, UNSPECIFIED, NamedSharding, GSPMDSharding, PositionalSharding,
     SingleDeviceSharding, parse_flatten_op_sharding)
 import jax._src.pjit as pjit_lib
+from jax._src.maps import xmap
 from jax._src.pjit import pjit, pjit_p
 from jax._src import mesh as mesh_lib
 from jax._src.interpreters import pxla

tests/python_callback_test.py

@@ -27,13 +27,13 @@
 from jax._src import core
 from jax._src import dispatch
 from jax._src import maps
+from jax._src.maps import xmap
 from jax._src import test_util as jtu
 from jax._src import util
 from jax._src.lib import xla_client
 from jax._src.lib import xla_extension_version
 from jax.experimental import io_callback
 from jax.experimental import pjit
-from jax.experimental.maps import xmap
 from jax.experimental.shard_map import shard_map
 import jax.numpy as jnp
 from jax.sharding import Mesh