Preserve runtime reduction axes in AdaScale ONNX conversion

[cherry77-cloud (taking-lying-flat)]

The opset-18 ReduceMean and ReduceL2 converters discard axes when constant resolution fails. An axes graph input, or an initializer routed through Identity, therefore becomes a reduction over every dimension. A following broadcast division can hide the error because the final shape is unchanged.

Preserve unresolved axes in the converted input mapping and read them in forward. Keep conversion-time folding for direct constants, so their existing path does not gain a runtime tensor-to-list conversion. Missing optional inputs still use the existing default semantics. Also correct ReduceL2 with empty axes and noop_with_empty_axes=1 to return the elementwise magnitude, as required by its square/sum/sqrt definition (ONNX ReduceL2-18).

With x = [[1,2,3],[4,5,6]] and axes = [1], x / ReduceMean(x, axes) previously returns [0.285714, 0.571429, 0.857143] in the first row. It now returns [0.5, 1, 1.5], exactly matching ONNX Runtime. The initializer control remains exact. Graph-input and Identity axes also match ONNX Runtime for ReduceL2.

Validation: 275 checks passed.

• Local source tests through the real get_pt_block API: 98 CPU + 98 CUDA passed. Covers both operators, graph inputs, Identity inputs, initializer/Constant controls, positive/negative/multiple axes, keepdims, omitted/empty axes and no-op behavior, changing axes between calls, normalization input gradients, and converted MatMul weight gradients and one SGD update against native PyTorch. The initial 92-case matrix on unchanged source gives 51 failures and 41 passes.
• pytest TrainingExtensions/onnx/test/python/test_onnx2torch.py TrainingExtensions/onnx/test/python/test_model_converter.py -q -ra -k 'not round_trip_with_qwen': 70 passed, 7 deselected. Those seven tests require Qwen/Qwen2.5-0.5B/backbone/model_sldynamic_cl32.onnx, which is absent locally; running them on both original and patched source produces the same FileNotFoundError before conversion.
• pytest TrainingExtensions/onnx/test/python/test_adascale.py::TestRequiredExtraBlockInputs TrainingExtensions/onnx/test/python/test_adascale.py::TestUpcastFp16BlockToBf16 TrainingExtensions/onnx/test/python/test_adascale.py::TestAdascaleQuantizer::test_block_level_api -q -ra: 9 passed.
• Repository pre-commit hooks and git diff --check: passed. The pre-commit interpreter setting was adapted locally to Python 3.12; hook definitions were unchanged.

Environment: Python 3.12, PyTorch 2.13.0+cu130, onnx2torch 1.5.15, RTX A1000 Laptop GPU. Tests use generated small graphs and source AIMET code; this does not claim a full pretrained-model PTQ evaluation. Additional validation scripts are outside the production-only diff.

Open PR searches for ReduceMean, ReduceL2, dynamic axes, and axes AdaScale found no duplicate. Implementation and validation were assisted by Codex.

Minimal graph-input reproducer

import numpy as np
import onnx_ir
import torch
from onnx import TensorProto, helper
from aimet_onnx.adascale.model_converter import get_pt_block

model = helper.make_model(
    helper.make_graph(
        [helper.make_node("ReduceMean", ["x", "axes"], ["mean"], keepdims=1),
         helper.make_node("Div", ["x", "mean"], ["y"])],
        "normalize",
        [helper.make_tensor_value_info("x", TensorProto.FLOAT, [2, 3]),
         helper.make_tensor_value_info("axes", TensorProto.INT64, [1])],
        [helper.make_tensor_value_info("y", TensorProto.FLOAT, [2, 3])],
    ),
    opset_imports=[helper.make_opsetid("", 18)],
)
model.ir_version = 10
block, _ = get_pt_block(onnx_ir.from_proto(model), (["x", "axes"], ["y"]))
x = torch.tensor([[1., 2., 3.], [4., 5., 6.]])
actual = block(x, torch.tensor([1]))
torch.testing.assert_close(actual, x / x.mean(dim=1, keepdim=True), rtol=0, atol=0)

[Abhi Khobare (abhikhobare)]

Thanks cherry77-cloud (@taking-lying-flat) for your contribution.

Your commit has been merged into develop now - here db50a5c with some additional changes.

Closing this PR now.