code: def testDatasetPartialBatchWithMixedOutputs(self, distribution):
      # Dynamic output size with a mix of static and dynamic outputs
      dataset = get_dataset_from_tensor_slices([5.0]).batch(2)
      input_iterator = iter(
          distribution.experimental_distribute_dataset(dataset)
      )
  
      @def_function.function
      def run(iterator):
          def computation(x):
              # Fixed size output with a dynamic sized output.
              return array_ops.zeros([3]), math_ops.square(x)
  
          return distribution.run(computation, args=(next(iterator),))
  
      results = run(input_iterator)
  
      # First result is fixed for all replicas.
      for replica_id in range(distribution.num_replicas_in_sync):
          self.assertAllEqual(
              [0.0, 0.0, 0.0],
              distribution.experimental_local_results(results[0])[
                  replica_id
              ],
          )
      # Only first replica has distributed dataset computation.
      self.assertAllEqual(
          [25.0], distribution.experimental_local_results(results[1])[0]
      )
      # Other replicas have no distributed dataset computation.
      for replica_id in range(1, distribution.num_replicas_in_sync):
          self.assertAllEqual(
              [],
              distribution.experimental_local_results(results[1])[
                  replica_id
              ],
          )
  
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 34 for calling the InputIterationTest obj's underlying member method code with example usage: obj.testDatasetPartialBatchWithMixedOutputs(distribution) and returns: distribution, array_ops,  math_ops