class_docstr: Neural Architecture Search (NAS) recurrent network cell.
  
  This implements the recurrent cell from the paper:
  
    https://arxiv.org/abs/1611.01578
  
  Barret Zoph and Quoc V. Le.
  "Neural Architecture Search with Reinforcement Learning" Proc. ICLR 2017.
  
  The class uses an optional projection layer.
  
  Example:
  
  >>> inputs = np.random.random([30,23,9]).astype(np.float32)
  >>> NASCell = tfa.rnn.NASCell(4)
  >>> rnn = tf.keras.layers.RNN(NASCell, return_sequences=True, return_state=True)
  >>> outputs, memory_state, carry_state = rnn(inputs)
  >>> outputs.shape
  TensorShape([30, 23, 4])
  >>> memory_state.shape
  TensorShape([30, 4])
  >>> carry_state.shape
  TensorShape([30, 4])
class_docstr_tok: ["Neural","Architecture","Search","(","NAS",")","recurrent","network","cell",".","This","implements","the","recurrent","cell","from","the","paper",":","https",":","\/\/arxiv.org\/abs\/1611.01578","Barret","Zoph","and","Quoc","V.","Le",".","``","Neural","Architecture","Search","with","Reinforcement","Learning","''","Proc",".","ICLR","2017",".","The","class","uses","an","optional","projection","layer",".","Example",":",">",">",">","inputs","=","np.random.random","(","[","30,23,9","]",")",".astype","(","np.float32",")",">",">",">","NASCell","=","tfa.rnn.NASCell","(","4",")",">",">",">","rnn","=","tf.keras.layers.RNN","(","NASCell",",","return_sequences=True",",","return_state=True",")",">",">",">","outputs",",","memory_state",",","carry_state","=","rnn","(","inputs",")",">",">",">","outputs.shape","TensorShape","(","[","30",",","23",",","4","]",")",">",">",">","memory_state.shape","TensorShape","(","[","30",",","4","]",")",">",">",">","carry_state.shape","TensorShape","(","[","30",",","4","]",")"]
docstr: Run one step of NAS Cell.
  
  Args:
    inputs: input Tensor, 2D, batch x num_units.
    state: This must be a list of state Tensors, both `2-D`, with column
      sizes `c_state` and `m_state`.
  
  Returns:
    A tuple containing:
    - A `2-D, [batch x output_dim]`, Tensor representing the output of
      the NAS Cell after reading `inputs` when previous state was
      `state`.
      Here output_dim is:
         projection if projection was set, units otherwise.
    - Tensor(s) representing the new state of NAS Cell after reading
      `inputs` when the previous state was `state`.  Same type and
      shape(s) as `state`.
  
  Raises:
    ValueError: If input size cannot be inferred from inputs via
      static shape inference.
docstr_tok: ["Run","one","step","of","NAS","Cell",".","Args",":","inputs",":","input","Tensor",",","2D",",","batch","x","num_units",".","state",":","This","must","be","a","list","of","state","Tensors",",","both","`","2-D","`",",","with","column","sizes","`","c_state","`","and","`","m_state","`",".","Returns",":","A","tuple","containing",":","-","A","`","2-D",",","[","batch","x","output_dim","]","`",",","Tensor","representing","the","output","of","the","NAS","Cell","after","reading","`","inputs","`","when","previous","state","was","`","state","`",".","Here","output_dim","is",":","projection","if","projection","was","set",",","units","otherwise",".","-","Tensor","(","s",")","representing","the","new","state","of","NAS","Cell","after","reading","`","inputs","`","when","the","previous","state","was","`","state","`",".","Same","type","and","shape","(","s",")","as","`","state","`",".","Raises",":","ValueError",":","If","input","size","can","not","be","inferred","from","inputs","via","static","shape","inference","."]
code: def call(self, inputs, state):
      """Run one step of NAS Cell.
  
      Args:
        inputs: input Tensor, 2D, batch x num_units.
        state: This must be a list of state Tensors, both `2-D`, with column
          sizes `c_state` and `m_state`.
  
      Returns:
        A tuple containing:
        - A `2-D, [batch x output_dim]`, Tensor representing the output of
          the NAS Cell after reading `inputs` when previous state was
          `state`.
          Here output_dim is:
             projection if projection was set, units otherwise.
        - Tensor(s) representing the new state of NAS Cell after reading
          `inputs` when the previous state was `state`.  Same type and
          shape(s) as `state`.
  
      Raises:
        ValueError: If input size cannot be inferred from inputs via
          static shape inference.
      """
      sigmoid = tf.math.sigmoid
      tanh = tf.math.tanh
      relu = tf.nn.relu
  
      c_prev, m_prev = state
  
      m_matrix = tf.matmul(m_prev, self.recurrent_kernel)
      inputs_matrix = tf.matmul(inputs, self.kernel)
  
      if self.use_bias:
          m_matrix = tf.nn.bias_add(m_matrix, self.bias)
  
      # The NAS cell branches into 8 different splits for both the hidden
      # state and the input
      m_matrix_splits = tf.split(
          axis=1, num_or_size_splits=self._NAS_BASE, value=m_matrix
      )
      inputs_matrix_splits = tf.split(
          axis=1, num_or_size_splits=self._NAS_BASE, value=inputs_matrix
      )
  
      # First layer
      layer1_0 = sigmoid(inputs_matrix_splits[0] + m_matrix_splits[0])
      layer1_1 = relu(inputs_matrix_splits[1] + m_matrix_splits[1])
      layer1_2 = sigmoid(inputs_matrix_splits[2] + m_matrix_splits[2])
      layer1_3 = relu(inputs_matrix_splits[3] * m_matrix_splits[3])
      layer1_4 = tanh(inputs_matrix_splits[4] + m_matrix_splits[4])
      layer1_5 = sigmoid(inputs_matrix_splits[5] + m_matrix_splits[5])
      layer1_6 = tanh(inputs_matrix_splits[6] + m_matrix_splits[6])
      layer1_7 = sigmoid(inputs_matrix_splits[7] + m_matrix_splits[7])
  
      # Second layer
      l2_0 = tanh(layer1_0 * layer1_1)
      l2_1 = tanh(layer1_2 + layer1_3)
      l2_2 = tanh(layer1_4 * layer1_5)
      l2_3 = sigmoid(layer1_6 + layer1_7)
  
      # Inject the cell
      l2_0 = tanh(l2_0 + c_prev)
  
      # Third layer
      l3_0_pre = l2_0 * l2_1
      new_c = l3_0_pre  # create new cell
      l3_0 = l3_0_pre
      l3_1 = tanh(l2_2 + l2_3)
  
      # Final layer
      new_m = tanh(l3_0 * l3_1)
  
      # Projection layer if specified
      if self.projection is not None:
          new_m = tf.matmul(new_m, self.projection_weights)
  
      return new_m, [new_c, new_m]
  
file_path: addons/tensorflow_addons/rnn/nas_cell.py
imports: import tensorflow
  from typeguard import typechecked
  from tensorflow_addons.rnn.abstract_rnn_cell import AbstractRNNCell
  from tensorflow_addons.utils.types import FloatTensorLike, TensorLike, Initializer
  from typing import Optional
label_desc: Use image node_id 5 for calling the NASCell obj's underlying member method code with example usage: obj.call(inputs, state) and returns: new_m,  list