Make some AQT dataclasses to use keyword-only fields (1/N)

This cl introduces an temporary decorator that will be temporarily used during this migration. The eventual goal is to enforce kw_only=True in all dataclasses unless it's not feasible, aiming to make AQT less error-prune and improve readability.

PiperOrigin-RevId: 631132072

MaxText/layers/quantizations.py

@@ -91,7 +91,9 @@ def _get_quant_config(config):
  else:
  drhs_bits = 8
  drhs_accumulator_dtype = jnp.int32
- drhs_local_aqt = aqt_config.LocalAqt(config.quantization_local_shard_count)
+ drhs_local_aqt = aqt_config.LocalAqt(
+ contraction_axis_shard_count=config.quantization_local_shard_count
+ )
  return aqt_config.config_v3(
  fwd_bits=8,
  dlhs_bits=8,