class_docstr: Tests that inspect the HLO for the sharding annotations.
    
class_docstr_tok: ["Tests","that","inspect","the","HLO","for","the","sharding","annotations","."]
code: def test_pjit_basic(self, in_shardings="P", out_shardings="P"):
      # Ensure that we can distinguish the inputs and outputs by shape
      def f_jax(x):  # f32[10,20] -> f32[20,10]
          return jnp.sin(x.T)
  
      pjit_kwargs = {}
      if in_shardings != "missing":
          pjit_kwargs["in_shardings"] = (
              P(None, "x") if in_shardings == "P" else None
          )
      if out_shardings != "missing":
          pjit_kwargs["out_shardings"] = (
              P("x", None) if out_shardings == "P" else None
          )
      f_jax = pjit.pjit(f_jax, **pjit_kwargs)
  
      x_shape = (10, 20)
      x = np.arange(np.prod(x_shape), dtype=np.float32).reshape(x_shape)
  
      self.log_jax_hlo(f_jax, [x], num_partitions=2)
  
      @tf.function(autograph=False, jit_compile=True)
      def f_tf(x):
          f_converted = jax2tf.convert(f_jax)
          if jtu.test_device_matches(["tpu"]):
              return tf.compat.v1.tpu.rewrite(
                  f_converted,
                  [tf.convert_to_tensor(x)],
                  device_assignment=self.device_assignment(
                      computation_shape=[1, 1, 1, 2],
                  ),
              )[0]
          else:
              return f_converted(x)
  
      # Annotation count for the input
      count_in_P = 1 if in_shardings == "P" else 0
      if config.jax2tf_default_native_serialization.value:
          # With native serialization even unspecified in_shardings turn into replicated
          count_in_replicated = (
              1 if in_shardings in [None, "missing"] else 0
          )
      else:
          count_in_replicated = 1 if in_shardings is None else 0
      # Annotation count for the output
      count_out_P = 1 if out_shardings == "P" else 0
      count_out_replicated = 1 if out_shardings is None else 0
  
      self.check_sharding(
          jax2tf.convert(f_jax),
          [x],
          checks=[
              # The argument
              (
                  r"f32\[10,20\].*custom_call_target.*Sharding.*sharding.*devices=\[1,2\]",
                  count_in_P,
              ),
              (
                  r"f32\[10,20\].*custom_call_target.*Sharding.*sharding.*replicated",
                  count_in_replicated,
              ),
              # The result
              (
                  r"f32\[20,10\].*custom_call_target.*Sharding.*sharding.*devices=\[2,1\]",
                  count_out_P,
              ),
              (
                  r"f32\[20,10\].*custom_call_target.*Sharding.*sharding.*replicated",
                  count_out_replicated,
              ),
              # No other shardings
              (
                  r"custom_call_target.*Sharding",
                  count_in_P
                  + count_in_replicated
                  + count_out_P
                  + count_out_replicated,
              ),
          ],
      )
  
      res_jax = f_jax(x)
      res_tf = f_tf(x)
      self.assertAllClose(res_tf.numpy(), res_jax)
  
imports: from collections.abc import Sequence
  import contextlib
  from functools import partial
  import logging
  import math
  import os
  import re
  from typing import Any
  import unittest
  from absl.testing import absltest
  import jax
  from jax._src import compiler
  from jax._src import config
  from jax._src import maps
  from jax._src import test_util
  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
  from jax.sharding import PartitionSpec
  import jax.numpy
  import numpy
  import tensorflow
  from jax.experimental.jax2tf.tests import tf_test_util
label_desc: Use image node_id 7 for calling the ShardingTest obj's underlying member method code with example usage: obj.test_pjit_basic(in_shardings, out_shardings) and returns: jnp, tf, f_converted