8c6540d278883cd3c78b501c69ad67a99d8b34a4,sonnet/python/modules/basic.py,BatchApply,_build,#BatchApply#,878

Before Change

    // with equal shape for the leading `n_dims` dimensions. Typically this is
    // the first input.
    example_input = tf.convert_to_tensor(
        nest.flatten(args)[self._input_example_index])
    def _split_to_original_leading_dims(result):
      return split_leading_dim(result, example_input, self._n_dims)
    return nest.map(_split_to_original_leading_dims, results)

After Change

      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_iterable([args, kwargs])
    merged_flattened = [merge_leading_dims(inp, self._n_dims)
                        for inp in flattened]
    merged_args, merged_kwargs = nest.pack_iterable_as([args, kwargs],
                                                       merged_flattened)

    results = self._module(*merged_args, **merged_kwargs)

Instances