code: def test_sharding_and_layouts(self):
      mesh = jtu.create_global_mesh((4, 2), ("x", "y"))
      shape = (4, 8)
      np_inp = np.arange(math.prod(shape)).reshape(shape)
      s = NamedSharding(mesh, P("x", "y"))
  
      compiled = (
          jax.jit(lambda x: x.T, in_shardings=s, out_shardings=s)
          .lower(
              np_inp, _in_layouts=layout.AUTO, _out_layouts=layout.AUTO
          )
          .compile()
      )
      out = compiled(np_inp)
      self.assertTupleEqual(
          compiled._input_layouts()[0][0]._minor_to_major, (1, 0)
      )
      self.assertTupleEqual(
          compiled._output_layouts()._minor_to_major, (0, 1)
      )
      self.assertArraysEqual(out, np_inp.T)
      self.assertEqual(out.sharding, s)
  
imports: import math
  import os
  from absl.testing import absltest
  import numpy
  import jax
  from jax.sharding import NamedSharding, PartitionSpec
  from jax._src import config
  from jax._src import layout
  from jax._src import test_util
  from jax._src.util import safe_zip
  from jax._src import xla_bridge
  from jax._src.lib import xla_extension_version
label_desc: Use image node_id 5 for calling the LayoutTest obj's underlying member method code with example usage: obj.test_sharding_and_layouts() without return types