class_docstr: d_model, n_ctdx, d_attn, n_head are all standard Transformer hyperparams
  attn_dropout, resid_dropout handle dropout rates for different parameters
  init_scale: scalar which can be used to change all inits
  mp_comm: a comm for sharding
class_docstr_tok: ["d_model",",","n_ctdx",",","d_attn",",","n_head","are","all","standard","Transformer","hyperparams","attn_dropout",",","resid_dropout","handle","dropout","rates","for","different","parameters","init_scale",":","scalar","which","can","be","used","to","change","all","inits","mp_comm",":","a","comm","for","sharding"]
docstr: figures out how much recompute to do on past
docstr_tok: ["figures","out","how","much","recompute","to","do","on","past"]
code: def _get_past_keys_values(self, past: torch.Tensor):
      """
      figures out how much recompute to do on past
      """
      # This handles the past context
      if past is not None:
          # We are sampling, and we have some precomputed queries/keys
          #
          # Or, we evaluated a shared context and want to do forward passes on continuations
          # (not yet supported)
          assert past.dim() == 4, f"wrong past shape: {past.size()}"
          # [2, batch_n, past_ctx_n, model_dim]
          past_keys, past_values = past
          return past_keys, past_values
      else:
          return None, None
  
file_path: summarize-from-feedback/summarize_from_feedback/models/attention.py
imports: from typing import Tuple
  import numpy
  import torch
  import torch.nn
  from summarize_from_feedback.models.ops import Conv1D
  from summarize_from_feedback.utils import exact_div
  from summarize_from_feedback.utils.dist_utils import Comm
label_desc: Use image node_id 2 for calling the Attention obj's underlying member method code with example usage: obj._get_past_keys_values(past) and returns: past_keys,  past_values, None,  None