Improve masked_scatter implementation documentation and code clarity

[Copilot]

This PR addresses the "potential improvement" suggestion for the masked_scatter operation by enhancing documentation and code clarity while maintaining the correct and optimal implementation.

Analysis of the Suggested "Simple" Approach

The PR discussion suggested using self * (1-mask) + mask * source as a simpler alternative. However, this approach is fundamentally incorrect for PyTorch's masked_scatter operation because:

1. Sequential consumption: masked_scatter consumes values from source sequentially (flattened), not element-wise
2. Shape flexibility: source can have arbitrary shape and more elements than needed
3. Broadcasting mismatch: The simple approach requires same-shaped tensors

Example Demonstrating Why Simple Approach Fails

import torch

self = torch.tensor([1.0, 2.0, 3.0, 4.0])
mask = torch.tensor([False, True, False, True])
source = torch.tensor([10.0, 20.0])  # Only 2 elements for 4 positions

# Correct masked_scatter result: [1.0, 10.0, 3.0, 20.0]
result = torch.masked_scatter(self, mask, source)

# Simple approach fails with shape mismatch:
# self * (1-mask) + mask * source  # RuntimeError: size mismatch

Current Implementation is Already Optimal

The existing ScatterND-based implementation is correct and optimal for the general case:

• Handles sequential consumption from flattened source
• Works with arbitrary source shapes
• Minimal ONNX operations required
• Proper broadcasting behavior

Improvements Made

Since the algorithm cannot be simplified without breaking correctness, this PR focuses on code quality:

Enhanced Documentation

• Comprehensive docstring explaining the operation semantics
• Clear explanation of why simple approaches don't work
• Detailed parameter and behavior documentation

Improved Code Readability

• Better variable names (true_indices, flattened_source, source_values)
• Logical grouping of related operations
• Clear inline comments explaining each step

Performance Maintained

• Zero overhead: Same number of ONNX operations
• Identical functionality: No behavioral changes
• Future-proof: Prevents incorrect "optimizations"

The implementation now serves as both a correct solution and educational resource for understanding the complexity required for masked_scatter operations.

Fixes #2113.

---

💬 Share your feedback on Copilot coding agent for the chance to win a $200 gift card! Click here to start the survey.

[justinchuby]

Thank you