code: def test_unroll(
      self,
      core_fn: Callable[[], snt.RNNCore],
      input_shape: Tuple[int],
      dtype: tf.DType,
      dynamic: bool,
      replicator_fn: tf.distribute.Strategy,
  ):
      replicator = replicator_fn()
      with replicator.scope():
          core = core_fn()
  
      def step_fn():
          def forward():
              unroll = (
                  snt.dynamic_unroll if dynamic else snt.static_unroll
              )
              sequence = tf.ones((1,) + input_shape, dtype)
              state = core.initial_state(input_shape[0])
              return unroll(core, sequence, state)
  
          return replicator.run(forward)
  
      # TpuReplicator doesn't support pure eager mode.
      if isinstance(replicator, snt_replicator.TpuReplicator):
          step_fn = tf.function(step_fn)
  
      # TODO(b/132329316) Remove when `xla.compile` allows tf.device(TPU).
      with tf.device(None):
          out_sequence, final_state = step_fn()
  
      self.assertSameValuePerReplica(replicator, out_sequence)
      self.assertSameValuePerReplica(replicator, final_state)
  
file_path: sonnet/sonnet/src/conformance/distribute_test.py
imports: from typing import Callable, Tuple
  from absl.testing import parameterized
  import sonnet
  from sonnet.src import test_utils
  from sonnet.src.conformance import descriptors
  from sonnet.src.conformance import goldens
  from sonnet.src.distribute import replicator
  from sonnet.src.distribute import replicator_test_utils
  import tensorflow
label_desc: Use image node_id 3 for calling the TpuReplicatorTest obj's underlying member method code with example usage: obj.test_unroll(core_fn, input_shape, dtype, dynamic, replicator_fn) and returns: replicator, unroll