docs: def inverse_categorical(x, mu=255):
      """
      Invert categorical samples
      """
      flat = x.reshape((-1, x.shape[-1]))
      indices = np.argmax(flat, axis=1).astype(np.float32)
      indices = (indices / mu) - 0.5
      inverted = inverse_mu_law(indices, mu=mu).reshape(x.shape[:-1])
      return ArrayWithUnits(inverted, x.dimensions[:2])