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: * `grad_output` is the gradient with respect to the output of about `forward` function
  
  This accumulates the gradients in the shared memory tensor and return the
  gradients with respect to the `last` result in the stack.
docstr_tok: ["*","`","grad_output","`","is","the","gradient","with","respect","to","the","output","of","about","`","forward","`","function","This","accumulates","the","gradients","in","the","shared","memory","tensor","and","return","the","gradients","with","respect","to","the","`","last","`","result","in","the","stack","."]
code: def backward(ctx, grad_output):
      """
      * `grad_output` is the gradient with respect to the output of about `forward` function
  
      This accumulates the gradients in the shared memory tensor and return the
      gradients with respect to the `last` result in the stack.
      """
      # Get the current size of the stack
      n = ctx._n
      # Get the accumulated gradients
      memory_grad = ctx._mem_grad
      # Add the gradients
      memory_grad[: n + 1] += grad_output
      # Return the gradients w.r.t to last value in the stack
      return None, None, memory_grad[n], None
  
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 StackFunction obj's underlying member method code with example usage: obj.backward(ctx, grad_output) and returns: None,  None,  memory_grad,  None