QuantizedMatmul

Matrix multiplication with a quantized weight matrix.

Abstract Signature:

QuantizedMatmul(x: Tensor, w: Tensor, scales: Tensor, biases: Tensor | None, transpose: bool = True, group_size: int | None, bits: int | None, mode: str = affine)

PyTorchJAX (Core)KerasTensorFlowApple MLXFlax NNXPaxML / Praxis

PyTorch

API: —

Strategy: Plugin (quantized_matmul)

JAX (Core)

API: —

Strategy: Plugin (quantized_matmul)

Keras

API: —

Strategy: Plugin (quantized_matmul)

TensorFlow

API: —

Strategy: Plugin (quantized_matmul)

Apple MLX

API: mlx.core.quantized_matmul

Strategy: Direct Mapping

Official Docs ↗

Flax NNX

API: —

Strategy: Plugin (quantized_matmul)

PaxML / Praxis

API: —

Strategy: Plugin (quantized_matmul)