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: Flattens and caches the tensor's batch_dims.
docstr_tok: ["Flattens","and","caches","the","tensor","'s","batch_dims","."]
code: def flatten(self, tensor):
      """Flattens and caches the tensor's batch_dims."""
      with tf.name_scope("batch_flatten"):
          if self._batch_dims == 1:
              return tensor
  
          self._original_tensor_shape = composite.shape(tensor)
  
          if tensor.shape[self._batch_dims :].is_fully_defined():
              return composite.reshape(
                  tensor,
                  [-1] + tensor.shape[self._batch_dims :].as_list(),
              )
  
          reshaped = composite.reshape(
              tensor,
              tf.concat(
                  [[-1], composite.shape(tensor)[self._batch_dims :]],
                  axis=0,
              ),
          )
          # If the batch dimensions are all defined but the rest are undefined,
          # `reshaped` will have None as the first squashed dim since we are calling
          # tf.shape above. Since we know how many batch_dims we have, we can check
          # if all the elements we want to squash are defined, allowing us to
          # call ensure_shape to set the shape of the squashed dim. Note that this
          # is only implemented for tf.Tensor and not SparseTensors.
          if (
              isinstance(tensor, tf.Tensor)
              and tensor.shape[: self._batch_dims].is_fully_defined()
          ):
              return tf.ensure_shape(
                  reshaped,
                  [
                      np.prod(
                          tensor.shape[: self._batch_dims],
                          dtype=np.int64,
                      )
                  ]
                  + tensor.shape[self._batch_dims :],
              )
          return reshaped
  
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 2 for calling the BatchSquash obj's underlying member method code with example usage: obj.flatten(tensor) and returns: reshaped, tensor, composite, tf