docs: def categorical(x, mu=255, normalize=True):
      """
      Mu-law compress a block of audio samples, and convert them into a
      categorical distribution
      """
  
      if normalize:
          # normalize the signal
          mx = x.max()
          x = np.divide(x, mx, where=mx != 0)
  
      # mu law compression
      x = mu_law(x)
  
      # translate and scale to [0, 1]
      x = (x - x.min()) * 0.5
  
      # convert to the range [0, 255]
      x = (x * mu).astype(np.uint8)
  
      # create the array to house the categorical representation
      c = np.zeros((np.product(x.shape), mu + 1), dtype=np.uint8)
      c[np.arange(len(c)), x.flatten()] = 1
  
      return ArrayWithUnits(
          c.reshape(x.shape + (mu + 1,)),
          x.dimensions + (IdentityDimension(),))