class_docstr: ### Stack Module
  
  This uses the stack function defined above, and does the necessary initializations.
class_docstr_tok: ["#","#","#","Stack","Module","This","uses","the","stack","function","defined","above",",","and","does","the","necessary","initializations","."]
docstr: * `n` is the size of the stack
  * `value` is the tensor that needs to be added to the stack
docstr_tok: ["*","`","n","`","is","the","size","of","the","stack","*","`","value","`","is","the","tensor","that","needs","to","be","added","to","the","stack"]
code: def append(self, n: int, value: torch.Tensor):
      """
      * `n` is the size of the stack
      * `value` is the tensor that needs to be added to the stack
      """
  
      # You need to get (use) the stack after adding a value.
      # Otherwise this implementation fails
      assert (
          n == 0 or self.last_get_n == n - 1
      ), f"{n}, {self.last_get_n}"
  
      # Do this without gradients
      with torch.no_grad():
          # Initialize the shared memory tensor to keep the stack
          if (
              self.memory is None
              or self.memory.shape[1:] != value.shape
          ):
              # This should only happen when the stack is empty
              assert n == 0
              # Create a tensor for the stack
              self.memory = value.new_zeros(
                  self.max_len, *value.shape, requires_grad=False
              )
              # Create a tensor to accumulate the gradients
              self.memory_grad = value.new_zeros(
                  self.memory.shape, requires_grad=False
              )
          # The memory is already initialized but we are resetting the stack.
          #
          # This could have been another function like `reset`, but
          # we found this easier to use.
          elif n == 0:
              # Reset accumulated gradients
              self.memory_grad.fill_(0.0)
  
          # Set the value in the correct position of the stack
          self.memory.data[n] = value.detach()
          # Keep track of the stack (for debugging)
          self.n = n
  
      # Keep track of the last value added to the stack.
      # We need this to be passed on to `StackFunction` in order
      # to get the gradients propagated backwards.
      self.last = value
  
imports: import math
  from typing import Optional
  import torch
  from torch import nn
  from labml_helpers.module import Module
  from labml_nn.transformers.feed_forward import FeedForward
  from labml_nn.transformers.mha import PrepareForMultiHeadAttention
  from labml_nn.utils import clone_module_list
label_desc: Use image node_id 2 for calling the Stack obj's underlying member method code with example usage: obj.append(n, value) without return types