peft.simulate_quantize

peft.simulate_quantize#

gemma.peft.simulate_quantize(
x: jax.Array | Any,
method: gemma.peft._quantization_utils.QuantizationMethod | str,
axis_to_reduce: int | None = None,
) → jax.Array | Any[source]

Quantizes the given array.

In this API, we do not actually quantize tensors as the output is not stored using less bits but rather simulate quantization to enable quantization aware training.

NOTE: you can use this implementation to evaluate a checkpoint as if it was quantized.

Parameters:

  • x – The array to simulate_quantize.

  • method – The quantization method to use.

  • axis_to_reduce – The axis to reduce the array over.

Returns:

The simulate_quantized array.