class_docstr: Conducts NMF and holds the models and components 
class_docstr_tok: ["Conducts","NMF","and","holds","the","models","and","components"]
docstr: Prepares the activations tensor for NMF by reshaping it from four dimensions
  (batch, layer, neuron, position) down to two:
  ( neuron (and layer), position (and batch) ).
  
  Args:
      activations (tensor): activations tensors of shape (batch, layers, neurons, positions) and float values
      from_layer (int or None): Start value. Used to indicate a range of layers whose activations are to
          be processed
      to_layer (int or None): End value. Used to indicate a range of layers
      collect_activations_layer_nums (list of ints or None): A list of layer IDs. Used to indicate specific
          layers whose activations are to be processed
docstr_tok: ["Prepares","the","activations","tensor","for","NMF","by","reshaping","it","from","four","dimensions","(","batch",",","layer",",","neuron",",","position",")","down","to","two",":","(","neuron","(","and","layer",")",",","position","(","and","batch",")",")",".","Args",":","activations","(","tensor",")",":","activations","tensors","of","shape","(","batch",",","layers",",","neurons",",","positions",")","and","float","values","from_layer","(","int","or","None",")",":","Start","value",".","Used","to","indicate","a","range","of","layers","whose","activations","are","to","be","processed","to_layer","(","int","or","None",")",":","End","value",".","Used","to","indicate","a","range","of","layers","collect_activations_layer_nums","(","list","of","ints","or","None",")",":","A","list","of","layer","IDs",".","Used","to","indicate","specific","layers","whose","activations","are","to","be","processed"]
code: def reshape_activations(
      activations,
      from_layer: Optional[int] = None,
      to_layer: Optional[int] = None,
      collect_activations_layer_nums: Optional[List[int]] = None,
  ):
      """Prepares the activations tensor for NMF by reshaping it from four dimensions
      (batch, layer, neuron, position) down to two:
      ( neuron (and layer), position (and batch) ).
  
      Args:
          activations (tensor): activations tensors of shape (batch, layers, neurons, positions) and float values
          from_layer (int or None): Start value. Used to indicate a range of layers whose activations are to
              be processed
          to_layer (int or None): End value. Used to indicate a range of layers
          collect_activations_layer_nums (list of ints or None): A list of layer IDs. Used to indicate specific
              layers whose activations are to be processed
      """
  
      if len(activations.shape) != 4:
          raise ValueError(
              f"The 'activations' parameter should have four dimensions: "
              f"(batch, layers, neurons, positions). "
              f"Supplied dimensions: {activations.shape}",
              "activations",
          )
  
      if collect_activations_layer_nums is None:
          collect_activations_layer_nums = list(
              range(activations.shape[1])
          )
  
      layer_nums_to_row_ixs = {
          layer_num: i
          for i, layer_num in enumerate(collect_activations_layer_nums)
      }
  
      if from_layer is not None or to_layer is not None:
          from_layer = from_layer if from_layer is not None else 0
          to_layer = (
              to_layer if to_layer is not None else activations.shape[0]
          )
  
          if from_layer == to_layer:
              raise ValueError(
                  f"from_layer ({from_layer}) and to_layer ({to_layer}) cannot be the same value. "
                  "They must be apart by at least one to allow for a layer of activations."
              )
  
          if from_layer > to_layer:
              raise ValueError(
                  f"from_layer ({from_layer}) cannot be larger than to_layer ({to_layer})."
              )
  
          layer_nums = list(range(from_layer, to_layer))
      else:
          layer_nums = sorted(layer_nums_to_row_ixs.keys())
  
      if any([num not in layer_nums_to_row_ixs for num in layer_nums]):
          available = sorted(layer_nums_to_row_ixs.keys())
          raise ValueError(
              f"Not all layers between from_layer ({from_layer}) and to_layer ({to_layer}) "
              f"have recorded activations. Layers with recorded activations are: {available}"
          )
  
      row_ixs = [
          layer_nums_to_row_ixs[layer_num] for layer_num in layer_nums
      ]
      activation_rows = [activations[:, row_ix] for row_ix in row_ixs]
      # Merge 'layers' and 'neuron' dimensions. Sending activations down from
      # (batch, layer, neuron, position) to (batch, neuron, position)
  
      merged_act = np.concatenate(activation_rows, axis=1)
      # merged_act = np.stack(activation_rows, axis=1)
      # 'merged_act' is now (batch, neuron (and layer), position)
      merged_act = merged_act.swapaxes(0, 1)
      # 'merged_act' is now (neuron (and layer), batch, position)
      merged_act = merged_act.reshape(merged_act.shape[0], -1)
  
      return merged_act
  
imports: import os
  import json
  import ecco
  from IPython import display
  from ecco import util, lm_plots
  import random
  import matplotlib.pyplot
  import numpy
  import torch
  from torch.nn import functional
  from sklearn import decomposition
  from typing import Dict, Optional, List, Tuple, Union
  from ecco.util import strip_tokenizer_prefix, is_partial_token
label_desc: Use image node_id 2 for calling the NMF obj's underlying member method code with example usage: obj.reshape_activations(activations, from_layer, to_layer, collect_activations_layer_nums) and returns: merged_act