class_docstr: Tests that inspect the HLO for the sharding annotations.
    
class_docstr_tok: ["Tests","that","inspect","the","HLO","for","the","sharding","annotations","."]
docstr: Check the sharding in TF.
  
  Args:
    f_tf: the TF callable
    args_tf: the TF args
    checks: a list of tuples. The first element is a regular expression, the
      second element is an integer representing the expected number of
      occurrences of the regular expression in the TF HLO. As a special case,
      the second element can be the result of `self.GEQ(v)` to check that
      the number of occurrences is greater or equal to a value.
docstr_tok: ["Check","the","sharding","in","TF",".","Args",":","f_tf",":","the","TF","callable","args_tf",":","the","TF","args","checks",":","a","list","of","tuples",".","The","first","element","is","a","regular","expression",",","the","second","element","is","an","integer","representing","the","expected","number","of","occurrences","of","the","regular","expression","in","the","TF","HLO",".","As","a","special","case",",","the","second","element","can","be","the","result","of","`","self.GEQ","(","v",")","`","to","check","that","the","number","of","occurrences","is","greater","or","equal","to","a","value","."]
code: def check_sharding(self, f_tf, args_tf: Sequence[Any], *, checks=()):
      """Check the sharding in TF.
  
      Args:
        f_tf: the TF callable
        args_tf: the TF args
        checks: a list of tuples. The first element is a regular expression, the
          second element is an integer representing the expected number of
          occurrences of the regular expression in the TF HLO. As a special case,
          the second element can be the result of `self.GEQ(v)` to check that
          the number of occurrences is greater or equal to a value.
      """
      tf_hlo = self.tf_hlo(f_tf, args_tf)
      for check_re, expected_count in checks:
          count = len(re.findall(check_re, tf_hlo))
          if isinstance(expected_count, int):
              self.assertEqual(
                  count,
                  expected_count,
                  (
                      f"regular expression `{check_re}` expected to occur "
                      f"{expected_count} times but occurs {count} times in "
                      f"the TF HLO.\nThis is the TF HLO:\n{tf_hlo}"
                  ),
              )
          elif (
              isinstance(expected_count, tuple)
              and expected_count[0] == ">="
          ):
              self.assertGreaterEqual(
                  count,
                  expected_count[1],
                  (
                      f"regular expression `{check_re}` expected to occur "
                      f"at least {expected_count[1]} times but occurs {count} times in "
                      f"the TF HLO.\nThis is the TF HLO:\n{tf_hlo}"
                  ),
              )
          else:
              assert False
  
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 6 for calling the ShardingTest obj's underlying member method code with example usage: obj.check_sharding(f_tf, args_tf) without return types