`QuantizationMode` should be algebraic instead of scalar

[AppAppWorks]

Previously MLX supported affine quantization only, it was fine to have groupSize and bits as separate arguments of quantization methods. However, the newly supported mxfp4 only allows groupSize = 32 and bits = 4, it'd be better to confine quantization modes with a full sum type, e.g.

enum QuantizationMode: Equatable {
  case affine(groupSize: Int, bits: Int)
  case mxfp4
}

If, in the future, MLX adds support to more quantization modes (e.g. mxfpN), mxfp4 can be converted into a convenience method to maintain source compatibility, i.e.

enum QuantizationMode: Equatable {
  case affine(groupSize: Int, bits: Int)
  case mxfpN(n: Int) 
}

extension QuantizationMode {
  static let mxfp4: Self = .mxfpN(n: 4)
}

mode == .mxfp4 // still works because of the Equatable conformance

It will be a source breaking change for all use sites of QuantizationMode though, but I believe the use of sum type in this case will make the API more extensible and fool-proof.

[davidkoski]

Interesting idea. We also use this enum in decoding the configuration -- we would need something to handle that as well.

Perhaps we could make new methods/types and deprecate the old ones to make it non-source breaking. From the recent change to add the mode I know this touches a lot of code (but luckily code that is pretty localized in its use).

[AppAppWorks]

After some more thoughts, this change will also necessitate changes in Quantized,

public protocol Quantized: Module {
    var groupSize: Int { get } // to be deprecated
    var bits: Int { get } // to be deprecated
    var mode: QuantizationMode { get }
}

This new, context-rich QuantizationMode shall also abstract away the computation of the quantization scale factor from Quantized.

extension QuantizationMode {
  var quantizationScale: Int {
    switch self {
    case .affine(_, let bits): 32 / bits
    case .mxfp4: 32 / 4
    // future peculiar modes can do their specific calculations here, oblivious to the quantized modules
    }
  }
}

class QuantizedLinear {
  open override var shape: (Int, Int) {
   let shape = self.shape
   return (shape.0, shape.1 * self.mode.quantizationScale)
  }
}