docstr: Tests that RNG with dist variables can be used as tf.function's arg.
docstr_tok: ["Tests","that","RNG","with","dist","variables","can","be","used","as","tf.function","'s","arg","."]
code: def testDistVarAsTFFunArg(self, strat, jit_replica_fn):
      """Tests that RNG with dist variables can be used as tf.function's arg."""
      strat_name = type(strat).__name__
      if "CentralStorage" in strat_name:
          self.skipTest(
              "CentralStorageStrategy wraps variable updates in merge_call which "
              "can't be called inside a tf.function that doesn't cover the entire "
              "replica function (the function passed to strategy.run)."
          )
      if "TPU" in strat_name and not jit_replica_fn:
          self.skipTest(
              "TPUStrategy requires the replica function (the function passed to "
              "strategy.run) to be decorated with tf.function"
          )
      coord = None
      if "ParameterServer" in strat_name:
          coord = coordinator_lib.ClusterCoordinator(strat)
      shape = [3, 4]
      dtype = dtypes.int32
      with strat.scope():
          gen = rng.Generator.from_seed(1234)
  
          @def_function.function
          def f(gen):  # the main focus
              t1 = gen.uniform_full_int(shape=shape, dtype=dtype)
              t2 = gen.uniform_full_int(shape=shape, dtype=dtype)
              t = array_ops_stack.stack([t1, t2])
              return t
  
          def g():
              return f(gen)
  
          replica_fn = def_function.function(g) if jit_replica_fn else g
          for _ in range(2):
              results = run_on_strategy(replica_fn, strat, coord)
              values = strat.experimental_local_results(results)
              n = get_num_local_replicas(strat, values)
              self.assertAllEqual(n, len(values))
              self.assertAllDifferent(values)
  
imports: import functools
  import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.compat import v2_compat
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import multi_process_runner
  from tensorflow.python.distribute import sharded_variable
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import values
  from tensorflow.python.distribute.coordinator import cluster_coordinator
  from tensorflow.python.distribute.mirrored_strategy import MirroredStrategy
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import test_combinations
  from tensorflow.python.framework import test_util
  from tensorflow.python.module import module
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import stateful_random_ops
  from tensorflow.python.platform import test
  from tensorflow.python.saved_model import load
  from tensorflow.python.saved_model import save
  from tensorflow.python.util import deprecation
label_desc: Use image node_id 8 for calling the GeneratorTest obj's underlying member method code with example usage: obj.testDistVarAsTFFunArg(strat, jit_replica_fn) and returns: t, f