code: def testMockWithSharding(self):
      num_shards = 16
      jax.config.update("mock_num_gpus", num_shards)
      mesh = jtu.create_global_mesh((num_shards,), ("x",))
  
      @partial(
          jax.jit,
          in_shardings=NamedSharding(
              mesh,
              P(
                  "x",
              ),
          ),
          out_shardings=NamedSharding(
              mesh,
              P(
                  "x",
              ),
          ),
      )
      def f(x, y):
          z = x @ y
          return z @ y
  
      shape = (64, 64)
      x = jnp.arange(math.prod(shape)).reshape(shape).astype(np.float32)
      y = x + 1
      f_lowered = f.lower(x, y)
      hlo = f_lowered.compiler_ir()
      self.assertIn('sharding = "{devices=[16,1]<=[16]}"', str(hlo))
  
file_path: jax/tests/mock_gpu_test.py
imports: from functools import partial
  import math
  from absl.testing import absltest
  import jax
  from jax import config
  from jax._src import test_util
  import jax.numpy
  from jax.sharding import NamedSharding
  from jax.sharding import PartitionSpec
  import numpy
label_desc: Use image node_id 3 for calling the MockGPUTest obj's underlying member method code with example usage: obj.testMockWithSharding() and returns: unknown