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: Create two tied ops to flatten and unflatten the front dimensions.
  
  Args:
    batch_dims: Number of batch dimensions the flatten/unflatten ops should
      handle.
  
  Raises:
    ValueError: if batch dims is negative.
docstr_tok: ["Create","two","tied","ops","to","flatten","and","unflatten","the","front","dimensions",".","Args",":","batch_dims",":","Number","of","batch","dimensions","the","flatten\/unflatten","ops","should","handle",".","Raises",":","ValueError",":","if","batch","dims","is","negative","."]
code: def __init__(self, batch_dims):
      """Create two tied ops to flatten and unflatten the front dimensions.
  
      Args:
        batch_dims: Number of batch dimensions the flatten/unflatten ops should
          handle.
  
      Raises:
        ValueError: if batch dims is negative.
      """
      if batch_dims < 0:
          raise ValueError("Batch dims must be non-negative.")
      self._batch_dims = batch_dims
      self._original_tensor_shape = None
  
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 1 to create a new BatchSquash object from inherited base classes: object with example: obj = BatchSquash(batch_dims)