def zero_state(self, batch_size, dtype):
    """Initialize the tuple of state values to zeros except write strengths.

    Args:
      batch_size: The size of the current batch.
      dtype: The default datatype to initialize to.

    Returns:
      A new NeuralStackState tuple.
    """
    parent_state = NeuralStackState(*super(NeuralStackCell, self).zero_state(
        batch_size, dtype))
    return NeuralStackState(
        controller_state=parent_state.controller_state,
        read_values=parent_state.read_values,
        memory_values=parent_state.memory_values,
        read_strengths=parent_state.read_strengths,
        write_strengths=self.initialize_write_strengths(batch_size))
Initialize the tuple of state values to zeros except write strengths.

    Args:
      batch_size: The size of the current batch.
      dtype: The default datatype to initialize to.

    Returns:
      A new NeuralStackState tuple.
    
tensor2tensor/models/research/neural_stack.py