[Qiming Bao (14H034160212)]

torch.softmax(x, dim, dtype=...) and torch.log_softmax(x, dim, dtype=...) cast the input to dtype before the operation. aten_softmax and aten_special_log_softmax (which aten::log_softmax.int maps to) cast the result instead, so the ONNX graph was Softmax(x) -> Cast(dtype):

• Integer input (torch.softmax(torch.tensor([[1, 2, 3]]), -1, dtype=torch.float32)): ONNX Softmax/LogSoftmax take floats only. log_softmax exports an invalid model (Type 'tensor(int64)' of input parameter ... of operator (LogSoftmax) ... is invalid), and softmax returned [0, 0, 0] instead of [0.090, 0.245, 0.665].
• float16 input with dtype=torch.float32, as Hugging Face models use for attention probabilities (softmax(attn_weights, dim=-1, dtype=torch.float32)): the softmax is computed in float16 and only the result is cast. ONNX Runtime's CPU provider hides this because it has no float16 Softmax and upcasts on its own, but providers with a float16 kernel compute it in float16. With optimize=False the exported graph is Softmax(float16) -> Cast.

This PR casts self to dtype before Softmax/LogSoftmax, and widens the input type constraint from TFloat to TReal, so integer inputs are accepted.

Tests

• New test_softmax_dtype_casts_the_input_first (softmax and log_softmax) in e2e_ops_tests.py, exported with optimize=False. It checks that an int64 input with dtype=torch.float32 matches PyTorch, and that for a float16 input the Softmax/LogSoftmax node is fed by a Cast(to=FLOAT). Both cases fail on main and pass with this change.
• The OpInfo tests for softmax and log_softmax now also run on integer dtypes: 12 passed, against 8 on main.
• e2e_ops_tests.py has the same 5 failures as on main in my environment (onnx_ir/onnx versions), plus the 2 new passing tests.
• ruff check and ruff format are clean.