Internal prototype

PiperOrigin-RevId: 645409887

Author: Cerebra Catalyst Team

aqt/jax/v2/aqt_quantizer.py

@@ -43,6 +43,7 @@ class Quantizer:
   _calibrator: AbstractAqtCalibration | None = utils.static_field(default=None)
   # Round up the calibration to power of 2 (po2).
   po2_scale: bool = utils.static_field()
+  scale_dtype: jnp.dtype | None = utils.static_field(default=None)
   # TODO(yichizh): Factor out auxilliary dataclasses into a separate file.
   context: utils.Context
 
@@ -85,6 +86,8 @@ def calibrate(self, x, *, calibration_axes) -> aqt_tensor.QTensor:
     bound = self._calibrator.get_bound(x, shared_axes, self.context)
     abs_max_mapped_to = self.numerics.abs_val_mapped_to()
     scale = bound / abs_max_mapped_to
+    if self.scale_dtype:
+      scale = scale.astype(self.scale_dtype)
 
     if self.po2_scale:
       # With floor the biggest value (we are using jnp.max) is in the range of

aqt/jax/v2/config.py

@@ -354,6 +354,24 @@ def get_numerics(bits):
   return cfg
 
 
+def set_scale_dtype(cfg: DotGeneral, scale_dtype: jnp.dtype):
+  """Set scale_dtype for dot_general config."""
+  assert isinstance(
+      cfg.fwd.dg_quantizer, aqt_dot_general.DefaultDotGeneralQuantizer
+  )
+  assert isinstance(
+      cfg.dlhs.dg_quantizer, aqt_dot_general.DefaultDotGeneralQuantizer
+  )
+  assert isinstance(
+      cfg.drhs.dg_quantizer, aqt_dot_general.DefaultDotGeneralQuantizer
+  )
+  cfg.fwd.dg_quantizer.lhs.scale_dtype = scale_dtype
+  cfg.fwd.dg_quantizer.rhs.scale_dtype = scale_dtype
+  cfg.dlhs.dg_quantizer.lhs.scale_dtype = scale_dtype
+  cfg.dlhs.dg_quantizer.rhs.scale_dtype = scale_dtype
+  cfg.drhs.dg_quantizer.lhs.scale_dtype = scale_dtype
+  cfg.drhs.dg_quantizer.rhs.scale_dtype = scale_dtype
+
 ################################################################################
 # Functions below are auxiliary config creators.