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,