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_eager_error(self, func="pjit_sharded"):
      if config.jax2tf_default_native_serialization.value:
          raise unittest.SkipTest(
              "There is no error in eager mode for native serialization"
          )
  
      # Define some test functions
      @partial(pjit.pjit, in_shardings=(P("x"),), out_shardings=None)
      def f_pjit_sharded(a):
          return a + a
  
      @partial(pjit.pjit, in_shardings=None, out_shardings=None)
      def f_pjit_replicated(a):
          return a + a
  
      def f_nested_pjit_sharded(a):
          return a + pjit.pjit(
              jnp.sin, in_shardings=(P("x"),), out_shardings=None
          )(a)
  
      def f_nested_pjit_replicated(a):
          return a + pjit.pjit(
              jnp.sin, in_shardings=None, out_shardings=None
          )(a)
  
      shape = (8, 10)
      a = np.arange(np.prod(shape), dtype=np.float32).reshape(shape)
  
      if func == "pjit_sharded":
          f_jax = f_pjit_sharded
      elif func == "pjit_replicated":
          f_jax = f_pjit_replicated
      elif func == "nested_pjit_sharded":
          f_jax = f_nested_pjit_sharded
      elif func == "nested_pjit_replicated":
          f_jax = f_nested_pjit_replicated
      else:
          assert False
  
      with Mesh(self.devices, axis_names=("x",)):
          _ = f_jax(a)
          with self.assertRaisesRegex(
              ValueError,
              "function with sharded arguments or results must be used under a `tf.function` context",
          ):
              jax2tf.convert(f_jax)(a)
  
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 13 for calling the ShardingTest obj's underlying member method code with example usage: obj.test_pjit_eager_error(func) and returns: unknown, unknown, unknown, unknown