class_docstr: Merges a number of leading dimensions of an input tensor to manipulate it.
  
  Merges a number of leading dimensions of a tensor into a single dimension,
  connects the provided module, then splits the leading dimension of the
  result to match the input.
  
  Input tensors whose rank is smaller than the number of dimensions to collapse
  (e.g. all scalar values, which are tensors of rank 0), are passed unaltered to
  the provided module.
  
  This is useful for applying some module to each timestep of a Time x Batch x N
  tensor. If a module is hard coded to only support 2D (Batch x N) then the
  full 3D Tensor cannot be provided. BatchApply will 'merge' the first two
  dimensions of the sequence tensor by reshaping to a (Time * Batch) x N Tensor,
  and then the internal module can be applied. The result of that operation is
  reshaped such that its first dimensions are split to match the leading
  dimensions of the input.
class_docstr_tok: ["Merges","a","number","of","leading","dimensions","of","an","input","tensor","to","manipulate","it",".","Merges","a","number","of","leading","dimensions","of","a","tensor","into","a","single","dimension",",","connects","the","provided","module",",","then","splits","the","leading","dimension","of","the","result","to","match","the","input",".","Input","tensors","whose","rank","is","smaller","than","the","number","of","dimensions","to","collapse","(","e.g",".","all","scalar","values",",","which","are","tensors","of","rank","0",")",",","are","passed","unaltered","to","the","provided","module",".","This","is","useful","for","applying","some","module","to","each","timestep","of","a","Time","x","Batch","x","N","tensor",".","If","a","module","is","hard","coded","to","only","support","2D","(","Batch","x","N",")","then","the","full","3D","Tensor","can","not","be","provided",".","BatchApply","will","'merge","'","the","first","two","dimensions","of","the","sequence","tensor","by","reshaping","to","a","(","Time","*","Batch",")","x","N","Tensor",",","and","then","the","internal","module","can","be","applied",".","The","result","of","that","operation","is","reshaped","such","that","its","first","dimensions","are","split","to","match","the","leading","dimensions","of","the","input","."]
code: def __call__(self, *args, **kwargs):
      example = first_leaf(args, kwargs)
      if example is None:
          raise ValueError(
              "BatchApply requires at least one tensor input."
          )
  
      num_dims = self.num_dims
      merge = lambda x: merge_leading_dims(x, num_dims=num_dims)
      split = lambda x: split_leading_dim(
          x, num_dims=num_dims, example=example
      )
  
      # Merge leading dimensions of inputs.
      # Example: [T, B, N] -> [T*B, N]
      args = tree.map_structure(merge, args)
      kwargs = tree.map_structure(merge, kwargs)
  
      # Compute merged output.
      # Example: [T*B, O]
      outputs = self.module(*args, **kwargs)
  
      # Split leading dimensions of output to match input.
      # Example: [T*B, O] -> [T, B, O]
      return tree.map_structure(split, outputs)
  
file_path: sonnet/sonnet/src/batch_apply.py
imports: from typing import Any, Callable, Optional, Sequence, Union
  import numpy
  from sonnet.src import base
  import tensorflow
  import tree
label_desc: Use image node_id 2 for calling the BatchApply obj's underlying member method code with example usage: obj.__call__() and returns: tree