code: def testMapFnWithDynamicInputs(self, distribution):
      def dataset_fn(_):
          data = array_ops.zeros((20, 300, 32), dtype=dtypes.int32)
          dataset = get_dataset_from_tensor_slices(data)
          dataset = dataset.batch(16)
          return dataset
  
      input_iterator = iter(
          distribution.distribute_datasets_from_function(dataset_fn)
      )
  
      def embedding_lookup(inputs):
          embedding_weights = array_ops.zeros((1, 128))
          flat_inputs = array_ops.reshape(inputs, [-1])
          embeddings = array_ops.gather(embedding_weights, flat_inputs)
          embeddings = array_ops.reshape(
              embeddings, inputs.shape.as_list() + [128]
          )
          return embeddings
  
      @def_function.function
      def step_fn(example):
          return map_fn.map_fn(
              embedding_lookup,
              example,
              fn_output_signature=dtypes.float32,
          )
  
      # This assumes that there are exactly 2 replicas
      outputs = distribution.experimental_local_results(
          distribution.run(step_fn, args=(next(input_iterator),))
      )
      self.assertAllEqual((16, 300, 32, 128), outputs[0].shape)
      self.assertAllEqual((4, 300, 32, 128), outputs[1].shape)
  
imports: from absl.testing import parameterized
  from tensorflow.python import tf2
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import device_util
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import reduce_util
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import test
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import control_flow_assert
  from tensorflow.python.ops import map_fn
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.ops.losses import losses
  from tensorflow.python.tpu import tpu
  from tensorflow.python.util import nest
label_desc: Use image node_id 30 for calling the InputIterationTest obj's underlying member method code with example usage: obj.testMapFnWithDynamicInputs(distribution) and returns: dataset, embeddings, map_fn