class_docstr: Facilitates flattening and unflattening batch dims of a tensor.
  
  Exposes a pair of matched faltten and unflatten methods. After flattening only
  1 batch dimension will be left. This facilitates evaluating networks that
  expect inputs to have only 1 batch dimension.
class_docstr_tok: ["Facilitates","flattening","and","unflattening","batch","dims","of","a","tensor",".","Exposes","a","pair","of","matched","faltten","and","unflatten","methods",".","After","flattening","only","1","batch","dimension","will","be","left",".","This","facilitates","evaluating","networks","that","expect","inputs","to","have","only","1","batch","dimension","."]
docstr: Unflattens the tensor's batch_dims using the cached shape.
docstr_tok: ["Unflattens","the","tensor","'s","batch_dims","using","the","cached","shape","."]
code: def unflatten(self, tensor):
      """Unflattens the tensor's batch_dims using the cached shape."""
      with tf.name_scope("batch_unflatten"):
          if self._batch_dims == 1:
              return tensor
  
          if self._original_tensor_shape is None:
              raise ValueError("Please call flatten before unflatten.")
  
          # pyformat: disable
          return composite.reshape(
              tensor,
              tf.concat(
                  [
                      self._original_tensor_shape[: self._batch_dims],
                      composite.shape(tensor)[1:],
                  ],
                  axis=0,
              ),
          )
  
file_path: agents/tf_agents/networks/utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import typing
  import numpy
  import tensorflow
  from tf_agents.keras_layers import permanent_variable_rate_dropout
  from tf_agents.typing import types
  from tf_agents.utils import composite
label_desc: Use image node_id 3 for calling the BatchSquash obj's underlying member method code with example usage: obj.unflatten(tensor) and returns: composite, tensor