code: def reshape_activations_to_3d(activations):
      """
      Reshape the activations tensors into a shape where it's easier to compare
      activation vectors.
      Args:
          activations: activations tensor of LM. Shape:
              (batch, layer, neuron, position)
      returns:
          activations: activations tensor reshaped into:
              (layer, neuron, batch + position)
      """
  
      # Swap axes from (0 batch, 1 layer, 2 neuron, 3 position)
      # to (0 layer, 1 neuron, 2 batch, 3 position)
      activations = np.moveaxis(activations, [0, 1, 2], [2, 0, 1])
      s = activations.shape
      acts = activations.reshape(s[0], s[1], -1)
      return acts
  
label_desc: Use image node_id 2 for calling a global function with example usage: reshape_activations_to_3d(activations) and returns: acts