Fix CrossEntropyLoss block to support multi-output models

[Rishi-Dave]

Summary

• artifacts.generate_artifacts(..., loss=LossType.CrossEntropyLoss) no longer aborts with i < node_->OutputDefs().size() when the base model has multi-dimensional outputs.
• The SoftmaxCrossEntropyLoss op produces two outputs (loss, log_prob); the second was being dropped by graph optimizers because it had no value_info entry, leaving the gradient builder to dereference a missing output def via O(1).

Motivation

Fixes #22465. Users hit a hard C++ assertion when training models like DistilBERT whose forward graph emits a multi-dimensional last-hidden-state tensor. The same pattern appears for any seq2seq / LM training setup that pipes a 3-D output into CrossEntropyLoss.

This is a Python-only change scoped to the onnxblock training-artifacts API; the core inference engine is unaffected.

Changes

• orttraining/orttraining/python/training/onnxblock/loss/loss.py — after appending the SoftmaxCrossEntropyLoss node, register a value_info entry for log_prob_output_name so its output def survives shape inference and graph cleanup. Idempotent — guarded against duplicate entries.
• orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py — new test_crossentropy_loss_multi_output_model builds a 3-D output toy model, calls generate_artifacts with LossType.CrossEntropyLoss, and asserts the saved training_model.onnx retains both outputs on the SCE node.

Test Plan

• New test exercises the previously-failing path:
  python -m pytest orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py::test_crossentropy_loss_multi_output_model -v
• Existing CE coverage:
  python -m pytest orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py -k crossentropy -v
• lintrunner clean on the diff.

Fixes #22465

[Copilot]

Pull request overview

Fixes onnxblock training-artifacts generation for LossType.CrossEntropyLoss when the base model output is multi-dimensional by ensuring the SoftmaxCrossEntropyLoss node’s second output (log_prob) is preserved through optimization/shape-inference.

Changes:

• Add a value_info entry for SoftmaxCrossEntropyLoss’s log_prob output to prevent it from being dropped during graph optimization.
• Add a regression test that exports a toy seq2seq-style (3-D output) model and verifies generate_artifacts succeeds and the saved training model retains both SCE outputs.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
orttraining/orttraining/python/training/onnxblock/loss/loss.py	Registers log_prob in value_info after adding SoftmaxCrossEntropyLoss to avoid optimizer pruning breaking gradient building.
orttraining/orttraining/test/python/orttraining_test_ort_apis_onnxblock.py	Adds regression test covering multi-dimensional model output + CrossEntropyLoss artifact generation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Rishi-Dave]

Thanks for the catch. Pushed 122e5bbcd5 (with prerequisite 1ee5f8248a) to address the regression around the unconditional get_output_from_output_name call.

Changes:

• Added _graph_utils.get_value_info_for_name(model, name) that resolves a ValueInfoProto by searching graph.output → graph.input → graph.value_info in order, raising LookupError only if none match.
• CrossEntropyLoss.build now uses the new helper at both call sites: the scores_elem_type lookup for the new log_prob value_info entry, and the earlier labels_input = copy.deepcopy(...) path. This restores the documented "output so far" behavior where loss_input_name may be an intermediate tensor with type info in value_info.
• The existing multi-output regression test still covers the original reported failure; both call sites now share a single resolution strategy.