code: def bin_laid_out_activations(layout, activations, grid_size, threshold=5):
      
  
      assert layout.shape[0] == activations.shape[0]
  
      # calculate which grid cells each activation's layout position falls into
      # first bin stays empty because nothing should be < 0, so we add an extra bin
      bins = np.linspace(0, 1, num=grid_size + 1)
      bins[-1] = np.inf  # last bin should include all higher values
      indices = np.digitize(layout, bins) - 1  # subtract 1 to account for empty first bin
  
      # because of thresholding we may need to return a variable number of means
      means, coordinates, counts = [], [], []
  
      # iterate over all grid cell coordinates to compute their average directions
      grid_coordinates = np.indices((grid_size, grid_size)).transpose().reshape(-1, 2)
      for xy_coordinates in grid_coordinates:
          mask = np.equal(xy_coordinates, indices).all(axis=1)
          count = np.count_nonzero(mask)
          if count > threshold:
              counts.append(count)
              coordinates.append(xy_coordinates)
              mean = np.average(activations[mask], axis=0)
              means.append(mean)
  
      assert len(means) == len(coordinates) == len(counts)
      if len(coordinates) == 0:
          raise RuntimeError("Binning activations led to 0 cells containing activations!")
  
      return means, coordinates, counts
docs: Given a layout and activations, overlays a grid on the layout and returns
      averaged activations for each grid cell. If a cell contains less than `threshold`
      activations it will be discarded, so the number of returned data is variable.