class_docstr: The `NestMap` network processes nested inputs via nested layers.
  
  It is a TF-Agents network that can be used to process nested inputs.
  
  Stateful Keras layers (e.g. LSTMCell, RNN, LSTM, TF-Agents DynamicUnroll)
  are all supported.  The `state_spec` of `NestMap` has a structure matching
  that of `nested_layers`.
  
  `NestMap` can be used in conjunction with `NestFlatten` and a combiner
  (e.g. `tf.keras.layers.Add` or `tf.keras.layers.Concatenate`) to process
  and aggregate in a preprocessing step.
  
  Usage:
  ```python
  net = NestMap({"inp1": layer1, "inp2": layer2})
  outputs, next_state = net({"inp1": inp1, "inp2": inp2}, state)
  ```
class_docstr_tok: ["The","`","NestMap","`","network","processes","nested","inputs","via","nested","layers",".","It","is","a","TF-Agents","network","that","can","be","used","to","process","nested","inputs",".","Stateful","Keras","layers","(","e.g",".","LSTMCell",",","RNN",",","LSTM",",","TF-Agents","DynamicUnroll",")","are","all","supported",".","The","`","state_spec","`","of","`","NestMap","`","has","a","structure","matching","that","of","`","nested_layers","`",".","`","NestMap","`","can","be","used","in","conjunction","with","`","NestFlatten","`","and","a","combiner","(","e.g",".","`","tf.keras.layers.Add","`","or","`","tf.keras.layers.Concatenate","`",")","to","process","and","aggregate","in","a","preprocessing","step",".","Usage",":","``","`","python","net","=","NestMap","(","{","``","inp1","''",":","layer1",",","``","inp2","''",":","layer2","}",")","outputs",",","next_state","=","net","(","{","``","inp1","''",":","inp1",",","``","inp2","''",":","inp2","}",",","state",")","``","`"]
code: def call(self, inputs, network_state=(), **kwargs):
      nest_utils.assert_same_structure(
          self._nested_layers,
          inputs,
          allow_shallow_nest1=True,
          message=(
              "`self.nested_layers` and `inputs` do not have matching structures"
          ),
      )
  
      if network_state:
          nest_utils.assert_same_structure(
              self.state_spec,
              network_state,
              allow_shallow_nest1=True,
              message="network_state and state_spec do not have matching structure",
          )
          nested_layers_state = network_state
      else:
          nested_layers_state = tf.nest.map_structure(
              lambda _: (), self._nested_layers
          )
  
      # Here we must use map_structure_up_to because nested_layers_state has a
      # "deeper" structure than self._nested_layers.  For example, an LSTM
      # layer's state is composed of a list with two tensors.  The
      # tf.nest.map_structure function would raise an error if two
      # "incompatible" structures are passed in this way.
      def _mapper(inp, layer, state):  # pylint: disable=invalid-name
          return layer(inp, network_state=state, **kwargs)
  
      outputs_and_next_state = nest_utils.map_structure_up_to(
          self._nested_layers,
          _mapper,
          inputs,
          self._nested_layers,
          nested_layers_state,
      )
  
      flat_outputs_and_next_state = nest_utils.flatten_up_to(
          self._nested_layers, outputs_and_next_state
      )
      flat_outputs, flat_next_state = zip(*flat_outputs_and_next_state)
  
      outputs = tf.nest.pack_sequence_as(
          self._nested_layers, flat_outputs
      )
      next_network_state = tf.nest.pack_sequence_as(
          self._nested_layers, flat_next_state
      )
  
      return outputs, next_network_state
  
file_path: agents/tf_agents/networks/nest_map.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  import typing
  import tensorflow.compat.v2
  from tf_agents.networks import network
  from tf_agents.networks import sequential
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 4 for calling the NestMap obj's underlying member method code with example usage: obj.call(inputs, network_state) and returns: outputs,  next_network_state, layer