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.)

def _build(self, *args, **kwargs):
    """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.)
    """
    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)