class_docstr: ### Stack Function implementation
  
  We implement a custom function instead of appending to a python list
  and then doing `torch.stack`.
  This greatly improves the performance over calling `torch.stack` at
  each step along the sequence.
  Everytime `torch.stack` is called, it creates a new tensor, while
  this method and the accompanying class `Stack` share memory for each step.
class_docstr_tok: ["#","#","#","Stack","Function","implementation","We","implement","a","custom","function","instead","of","appending","to","a","python","list","and","then","doing","`","torch.stack","`",".","This","greatly","improves","the","performance","over","calling","`","torch.stack","`","at","each","step","along","the","sequence",".","Everytime","`","torch.stack","`","is","called",",","it","creates","a","new","tensor",",","while","this","method","and","the","accompanying","class","`","Stack","`","share","memory","for","each","step","."]
docstr: * `ctx` is the context of the function (which lets us cache stuff)
  * `memory` is the shared memory tensor where we stack and store the values of each step (keys & values)
  * `memory_grad` is the shared memory tensor to store and accumulate gradients of each step
  * `last` is the last value stacked
  * `n` is the number of steps (i.e. size of the stack)
  
  This returns the stacked tensor for steps upto `n`.
docstr_tok: ["*","`","ctx","`","is","the","context","of","the","function","(","which","lets","us","cache","stuff",")","*","`","memory","`","is","the","shared","memory","tensor","where","we","stack","and","store","the","values","of","each","step","(","keys","&","values",")","*","`","memory_grad","`","is","the","shared","memory","tensor","to","store","and","accumulate","gradients","of","each","step","*","`","last","`","is","the","last","value","stacked","*","`","n","`","is","the","number","of","steps","(","i.e",".","size","of","the","stack",")","This","returns","the","stacked","tensor","for","steps","upto","`","n","`","."]
code: def forward(ctx, memory, memory_grad, last, n):
      """
      * `ctx` is the context of the function (which lets us cache stuff)
      * `memory` is the shared memory tensor where we stack and store the values of each step (keys & values)
      * `memory_grad` is the shared memory tensor to store and accumulate gradients of each step
      * `last` is the last value stacked
      * `n` is the number of steps (i.e. size of the stack)
  
      This returns the stacked tensor for steps upto `n`.
      """
  
      # Cache accumulated gradients
      ctx._mem_grad = memory_grad
      # Cache the size of the stack
      ctx._n = n
      # Return the stack
      return memory[: n + 1]
  
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 1 for calling the StackFunction obj's underlying member method code with example usage: obj.forward(ctx, memory, memory_grad, last, n) and returns: memory