class_docstr: Unroll-aware deep RNN module.
  
  Sonnet's DeepRNN steps through RNNCores sequentially which can result in a
  performance hit, in particular when using Transformers. This module adds an
  unroll() method which unrolls each module in the DeepRNN individually,
  allowing efficient implementation of the unroll operation. For example, a
  Transformer can 'unroll' by evaluating the whole sequence at once (this being
  one of the advantages of Transformers over e.g. LSTMs).
  
  Any RNNCore passed to this module should implement unroll(). Failure to so
  may cause the RNNCore not to be called properly. For example, passing a
  partial function application of a snt.RNNCore to this module will fail (this
  is also true for snt.DeepRNN). However, the special case of passing in a
  RNNCore object that does not implement unroll() is supported and will be
  dynamically unrolled. Implement unroll() to override this behavior with
  static unrolling.
  
  Stateless modules (i.e. anything other than an RNNCore) which do not
  implement unroll() are batch applied over the time and batch axes
  simultaneously. Effectively, this means that such modules may be applied to
  fairly large batches, potentially leading to out-of-memory issues.
class_docstr_tok: ["Unroll-aware","deep","RNN","module",".","Sonnet","'s","DeepRNN","steps","through","RNNCores","sequentially","which","can","result","in","a","performance","hit",",","in","particular","when","using","Transformers",".","This","module","adds","an","unroll","(",")","method","which","unrolls","each","module","in","the","DeepRNN","individually",",","allowing","efficient","implementation","of","the","unroll","operation",".","For","example",",","a","Transformer","can","'unroll","'","by","evaluating","the","whole","sequence","at","once","(","this","being","one","of","the","advantages","of","Transformers","over","e.g",".","LSTMs",")",".","Any","RNNCore","passed","to","this","module","should","implement","unroll","(",")",".","Failure","to","so","may","cause","the","RNNCore","not","to","be","called","properly",".","For","example",",","passing","a","partial","function","application","of","a","snt.RNNCore","to","this","module","will","fail","(","this","is","also","true","for","snt.DeepRNN",")",".","However",",","the","special","case","of","passing","in","a","RNNCore","object","that","does","not","implement","unroll","(",")","is","supported","and","will","be","dynamically","unrolled",".","Implement","unroll","(",")","to","override","this","behavior","with","static","unrolling",".","Stateless","modules","(","i.e",".","anything","other","than","an","RNNCore",")","which","do","not","implement","unroll","(",")","are","batch","applied","over","the","time","and","batch","axes","simultaneously",".","Effectively",",","this","means","that","such","modules","may","be","applied","to","fairly","large","batches",",","potentially","leading","to","out-of-memory","issues","."]
docstr: Unroll each layer individually.
  
  Calls unroll() on layers which support it, all other layers are
  batch-applied over the first two axes (assumed to be the time and batch
  axes).
  
  Args:
    inputs: A nest of `tf.Tensor` in time-major format.
    state: The RNN core state.
    sequence_length: How long the static_unroll should go for.
  
  Returns:
    Nested sequence output of RNN, and final state.
  
  Raises:
    ValueError if the length of `state` does not match the number of
    unrollable layers.
docstr_tok: ["Unroll","each","layer","individually",".","Calls","unroll","(",")","on","layers","which","support","it",",","all","other","layers","are","batch-applied","over","the","first","two","axes","(","assumed","to","be","the","time","and","batch","axes",")",".","Args",":","inputs",":","A","nest","of","`","tf.Tensor","`","in","time-major","format",".","state",":","The","RNN","core","state",".","sequence_length",":","How","long","the","static_unroll","should","go","for",".","Returns",":","Nested","sequence","output","of","RNN",",","and","final","state",".","Raises",":","ValueError","if","the","length","of","`","state","`","does","not","match","the","number","of","unrollable","layers","."]
code: def unroll(
      self,
      inputs: types.NestedTensor,
      state: base.State,
      sequence_length: int,
  ) -> Tuple[types.NestedTensor, base.State]:
      """Unroll each layer individually.
  
      Calls unroll() on layers which support it, all other layers are
      batch-applied over the first two axes (assumed to be the time and batch
      axes).
  
      Args:
        inputs: A nest of `tf.Tensor` in time-major format.
        state: The RNN core state.
        sequence_length: How long the static_unroll should go for.
  
      Returns:
        Nested sequence output of RNN, and final state.
  
      Raises:
        ValueError if the length of `state` does not match the number of
        unrollable layers.
      """
      if len(state) != self._num_unrollable:
          raise ValueError(
              "DeepRNN was called with the wrong number of states. The length of "
              "`state` does not match the number of unrollable layers."
          )
  
      states = iter(state)
      outputs = inputs
      next_states = []
      for layer in self._layers:
          if hasattr(layer, "unroll"):
              # The length of the `states` list was checked above.
              outputs, next_state = layer.unroll(
                  outputs, next(states), sequence_length
              )
              next_states.append(next_state)
          else:
              # Couldn't unroll(); assume that this is a stateless module.
              outputs = snt.BatchApply(layer, num_dims=2)(outputs)
  
      return outputs, tuple(next_states)
  
imports: import functools
  from typing import NamedTuple, Optional, Sequence, Tuple
  from absl import logging
  from acme import types
  from acme.tf import savers
  from acme.tf import utils
  from acme.tf.networks import base
  import sonnet
  import tensorflow
  import tensorflow_probability
  import tree
label_desc: Use image node_id 2 for calling the DeepRNN obj's underlying member method code with example usage: obj.unroll(inputs, state, sequence_length) and returns: outputs,  tuple