def expand_amax(model):
    """Expand per-tensor amax to be per channel, where each channel is assigned the per-tensor amax."""

    for name, mod in model.named_modules():
        if hasattr(mod, "_weight_quantizer") and mod._weight_quantizer.axis is not None:
            k = mod.weight.shape[0]
            amax = mod._weight_quantizer._amax.detach()
            mod._weight_quantizer._amax = torch.ones(k, dtype=amax.dtype, device=amax.device) * amax
            print(f"expanding {name} {amax} -> {mod._weight_quantizer._amax}")
Expand per-tensor amax to be per channel, where each channel is assigned the per-tensor amax.
dpkd/transformers/examples/research_projects/quantization-qdqbert/quant_trainer.py