class_docstr: Squash the outer dimensions of input tensors; unsquash outputs.
  
  This layer wraps a Keras layer `wrapped` that cannot handle more than one
  batch dimension.  It squashes inputs' outer dimensions to a single larger
  batch then unsquashes the outputs of `wrapped`.
  
  The outer dimensions are the leftmost `rank(inputs) - inner_rank` dimensions.
  
  Examples:
  
  ```python
  batch_norm = tf.keras.layers.BatchNormalization(axis=-1)
  layer = SquashedOuterWrapper(wrapped=batch_norm, inner_rank=3)
  
  inputs_0 = tf.random.normal((B, H, W, C))
  # batch_norm sees tensor of shape [B, H, W, C]
  # outputs_1 shape is [B, H, W, C]
  outputs_0 = layer(inputs_0)
  
  inputs_1 = tf.random.normal((B, T, H, W, C))
  # batch_norm sees a tensor of shape [B * T, H, W, C]
  # outputs_1 shape is [B, T, H, W, C]
  outputs_1 = layer(inputs_1)
  
  inputs_2 = tf.random.normal((B1, B2, T, H, W, C))
  # batch_norm sees a tensor of shape [B1 * B2 * T, H, W, C]
  # outputs_2 shape is [B1, B2, T, H, W, C]
  outputs_2 = layer(inputs_2)
  ```
class_docstr_tok: ["Squash","the","outer","dimensions","of","input","tensors",";","unsquash","outputs",".","This","layer","wraps","a","Keras","layer","`","wrapped","`","that","can","not","handle","more","than","one","batch","dimension",".","It","squashes","inputs","'","outer","dimensions","to","a","single","larger","batch","then","unsquashes","the","outputs","of","`","wrapped","`",".","The","outer","dimensions","are","the","leftmost","`","rank","(","inputs",")","-","inner_rank","`","dimensions",".","Examples",":","``","`","python","batch_norm","=","tf.keras.layers.BatchNormalization","(","axis=-1",")","layer","=","SquashedOuterWrapper","(","wrapped=batch_norm",",","inner_rank=3",")","inputs_0","=","tf.random.normal","(","(","B",",","H",",","W",",","C",")",")","#","batch_norm","sees","tensor","of","shape","[","B",",","H",",","W",",","C","]","#","outputs_1","shape","is","[","B",",","H",",","W",",","C","]","outputs_0","=","layer","(","inputs_0",")","inputs_1","=","tf.random.normal","(","(","B",",","T",",","H",",","W",",","C",")",")","#","batch_norm","sees","a","tensor","of","shape","[","B","*","T",",","H",",","W",",","C","]","#","outputs_1","shape","is","[","B",",","T",",","H",",","W",",","C","]","outputs_1","=","layer","(","inputs_1",")","inputs_2","=","tf.random.normal","(","(","B1",",","B2",",","T",",","H",",","W",",","C",")",")","#","batch_norm","sees","a","tensor","of","shape","[","B1","*","B2","*","T",",","H",",","W",",","C","]","#","outputs_2","shape","is","[","B1",",","B2",",","T",",","H",",","W",",","C","]","outputs_2","=","layer","(","inputs_2",")","``","`"]
docstr: Initialize `SquashedOuterWrapper`.
  
  Args:
    wrapped: The keras layer to wrap.
    inner_rank: The inner rank of inputs that will be passed to the layer.
      This value allows us to infer the outer batch dimension regardless of
      the input shape to `build` or `call`.
    **kwargs: Additional arguments for keras layer construction.
  
  Raises:
    ValueError: If `wrapped` has method `get_initial_state`, because
      we do not know how to handle the case of multiple inputs and
      the presence of this method typically means an RNN or RNN-like
      layer which accepts separate state tensors.
docstr_tok: ["Initialize","`","SquashedOuterWrapper","`",".","Args",":","wrapped",":","The","keras","layer","to","wrap",".","inner_rank",":","The","inner","rank","of","inputs","that","will","be","passed","to","the","layer",".","This","value","allows","us","to","infer","the","outer","batch","dimension","regardless","of","the","input","shape","to","`","build","`","or","`","call","`",".","*","*","kwargs",":","Additional","arguments","for","keras","layer","construction",".","Raises",":","ValueError",":","If","`","wrapped","`","has","method","`","get_initial_state","`",",","because","we","do","not","know","how","to","handle","the","case","of","multiple","inputs","and","the","presence","of","this","method","typically","means","an","RNN","or","RNN-like","layer","which","accepts","separate","state","tensors","."]
code: def __init__(
      self,
      wrapped: tf.keras.layers.Layer,
      inner_rank: int,
      **kwargs: Mapping[Text, Any],
  ):
      """Initialize `SquashedOuterWrapper`.
  
      Args:
        wrapped: The keras layer to wrap.
        inner_rank: The inner rank of inputs that will be passed to the layer.
          This value allows us to infer the outer batch dimension regardless of
          the input shape to `build` or `call`.
        **kwargs: Additional arguments for keras layer construction.
  
      Raises:
        ValueError: If `wrapped` has method `get_initial_state`, because
          we do not know how to handle the case of multiple inputs and
          the presence of this method typically means an RNN or RNN-like
          layer which accepts separate state tensors.
      """
      if getattr(wrapped, "get_initial_state", None) is not None:
          raise ValueError(
              "`wrapped` has method `get_initial_state`, which means its inputs "
              "will include separate state tensors.  This is not supported by "
              "`SquashedOuterWrapper`.  wrapped: {}".format(wrapped)
          )
      self._inner_rank = inner_rank
      self._wrapped = wrapped
      super(SquashedOuterWrapper, self).__init__(**kwargs)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Any, Mapping, Text
  import numpy
  import tensorflow
  from tf_agents.networks import utils
label_desc: Use image node_id 1 to create a new SquashedOuterWrapper object from inherited base classes: tf.keras.layers with example: obj = SquashedOuterWrapper(wrapped, inner_rank)