code: def _build(self, *args, **kwargs):
      
      flattened = nest.flatten([args, kwargs])
      merged_flattened = [
          merge_leading_dims(inp, self._n_dims) if inp is not None else None
          for inp in flattened]
      merged_args, merged_kwargs = nest.pack_sequence_as([args, kwargs],
                                                         merged_flattened)
  
      results = self._module(*merged_args, **merged_kwargs)
  
      # Unmerging takes the sizes of the leading dimensions from an input example
      # with equal shape for the leading `n_dims` dimensions. Typically this is
      # the first input.
      example_input = tf.convert_to_tensor(flattened[self._input_example_index])
      def _split_to_original_leading_dims(result):
        if result is None:
          return None
        else:
          return split_leading_dim(result, example_input, self._n_dims)
  
      flat_results = nest.flatten(results)
      flat_unmerged_results = [_split_to_original_leading_dims(result)
                               for result in flat_results]
      return nest.pack_sequence_as(results, flat_unmerged_results)
docs: Connects the BatchApply module into the graph.
  
      Args:
        *args: a Tensor or a nested list or dictionary of Tensors. The input
            tensors will have their first dimensions merged, then an op or a
            module will be called on the input. The first dimension of the output
            tensor(s) will be split again based on the leading dimensions of the
            first input tensor.
        **kwargs: Dictionary of named arguments; used in the same way as `*args`.
  
      Returns:
        A Tensor or nested list or dictionary of Tensors as a result of applying
        the process above. ("None" return values are also supported.)