split out call_jax changes

Author: tengyifei

test/run_tests.sh

@@ -211,7 +211,6 @@ function run_xla_op_tests2 {
   run_test "$CDIR/test_callback.py"
   XLA_USE_SPMD=1 run_test "$CDIR/test_callback.py"
   run_test "$CDIR/test_jax_interop.py"
-  run_test "$CDIR/test_jax_interop_spmd.py"
   run_test "$CDIR/test_assume_pure.py"
   run_test "$CDIR/test_assume_pure_spmd.py"
 }

test/test_jax_interop_spmd.py

@@ -38,7 +38,6 @@ python3 "$TEST_CDIR/scan/test_scan_layers.py"
 python3 "$TEST_CDIR/test_gru.py"
 python3 "$TEST_CDIR/test_assume_pure.py"
 python3 "$TEST_CDIR/test_assume_pure_spmd.py"
-python3 "$TEST_CDIR/test_jax_interop_spmd.py"
 python3 "$TEST_CDIR/test_as_stride_use_slice.py"
 run_xla_hlo_debug python3 "$TEST_CDIR/scan/test_scan_debug.py"
 python3 "$TEST_CDIR/test_pallas.py" -v

test/tpu/run_tests.sh

@@ -920,19 +920,8 @@ def get_xla_computation():
       import torch_xla.debug.profiler as xp
       # If we see this trace span in the profiler, we'll know that there's a cache miss.
       with xp.Trace('jax_to_xla_computation'):
-        jitted = jax.jit(fn, keep_unused=True)
-
-        def do_lower():
-          import torch_xla.runtime as xr
-          import torch_xla.distributed.spmd as xs
-          if xr.is_spmd():
-            mesh = xs.get_global_mesh()
-            if mesh is not None:
-              with mesh.get_jax_mesh():
-                return jitted.lower(*sample_tensor_args)
-          return jitted.lower(*sample_tensor_args)
-
-        hlo_ir = do_lower().compiler_ir('hlo')
+        lowered = jax.jit(fn, keep_unused=True).lower(*sample_tensor_args)
+        hlo_ir = lowered.compiler_ir('hlo')
         assert len(traced_out_spec) == 1, \
             "fn must be traced to obtain the output tree spec"
         spec = traced_out_spec[0]