[xuan7zhang]

Description

For wrapper subclass inputs, such as the AsyncCollectiveTensor output of
RowwiseParallel, the cross_entropy autograd handler now uses the
CompositeImplicitAutograd decomposition. Plain tensors and DTensors do not change.

I used Claude Code to help write this change, and I reviewed and tested it myself.

Related Issue

Fixes #1026

Motivation and Context

The handler called the kernel with _op_dk(PrivateUse1), which skips the Python key.
AsyncCollectiveTensor did not unwrap, and the kernel failed with
"Tensor not bound: lazy alloc kAlloc may not have executed".

No open PR covers this.

Type of Change

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation update
• Refactoring (no functional changes)

Testing

Eager, 4 devices (trn2.3xlarge). The DTensor tests use world_size = 8, so I set it
to 4 locally for these runs. That change is not in the PR.

Ran locally (no Neuron hardware needed):

pre-commit run --all-files    # passed

Ran on Neuron hardware:

python -m pytest tests/distributed/dtensor/test_cross_entropy_loss.py -q \
  -k "rowwise_parallel_output or replicated or batch_sharded or test_cross_entropy_loss_parallel"
# 5 passed, 1 xpassed
# without the fix, the new test fails with "Tensor not bound"

python -m pytest tests/python_ops/test_cross_entropy_loss.py -q -n 4
# 81 passed

Checklist

• Added or updated tests for this change
• I've performed a self-review for my code
• Ran the tests, or explained under Testing why I could not
• Updated the corresponding documentation
• My changes generate no new warnings
• Said under Description if an AI assistant helped with this change (see AI_POLICY.md)

[kunind-aws]

Thanks for the fix, @xuan7zhang, and for taking the time to send it. We have already fixed this internally, so we'll close this one. We appreciate it.