Implement aten::feature_dropout

[Copilot]

This PR implements the missing aten::feature_dropout operation in onnxscript/function_libs/torch_lib/ops/core.py.

What is feature dropout?

Feature dropout (also known as channel dropout or spatial dropout) is a regularization technique that differs from regular dropout by dropping entire feature maps/channels rather than individual elements:

• Regular dropout: For each element in the tensor, randomly set it to 0 with probability p
• Feature dropout: For each channel/feature map, randomly set the entire channel to 0 with probability p

This is particularly useful in convolutional neural networks as it encourages the network not to rely too heavily on specific feature detectors while maintaining spatial correlations within each feature map.

Implementation Details

The implementation:

1. Adds proper decorator: @torch_op("aten::feature_dropout", trace_only=True)
2. Handles multiple tensor dimensions:
   • 2D tensors [N, C]: creates mask of shape [N, C]
   • Higher dimensions [N, C, H, W, ...]: creates mask of shape [N, C, 1, 1, ...] for broadcasting
3. Uses ONNX operations throughout: No Python control flow to ensure proper ONNX graph generation
4. Implements correct semantics:
   • Inference mode (train=False): returns input unchanged
   • Zero dropout (p=0): returns input unchanged
   • Training mode: applies feature-wise dropout with proper scaling 1/(1-p)

Key ONNX Operations Used

• RandomUniformLike for generating random values
• Where for conditional logic
• ConstantOfShape for creating appropriately shaped tensors
• Shape, Size, Gather, Concat for dynamic shape handling

Validation

• ✅ All existing tests continue to pass (730 validation tests)
• ✅ Function properly imports and has correct signature
• ✅ Behavior verified against PyTorch's torch.ops.aten.feature_dropout reference implementation
• ✅ Handles edge cases (inference mode, zero dropout) correctly

Fixes #2403.

---

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

[justinchuby]

Use op.Dropout for it?

[Copilot]

Great suggestion! I've refactored the implementation to use op.Dropout instead of manual random generation. The new approach creates a dummy tensor with the appropriate feature dropout shape [N, C, 1, 1, ...], applies op.Dropout to it to get the scaled mask, then multiplies with the input. This is much cleaner and leverages the existing ONNX Dropout operation for all the logic (training mode, scaling, random generation). Commit: 29f1968

[justinchuby]

Do

if p == 0 or not train:
    return input

# Do this other thing

[Copilot]

Added early return conditions as requested. The function now returns input directly when p == 0 or not train, avoiding unnecessary computations. Commit: af8d622

[justinchuby]

Suggested change

	batch_size = op.Gather(input_shape, [0])
	channel_size = op.Gather(input_shape, [1])
	batch_size = op.Shape(input, start=0, end=0)
	batch_size = op.Shape(input, start=1, end=1)

[Copilot]

Updated to use op.Shape with start/end parameters for more efficient shape extraction. Fixed the typo in the suggestion where both lines assigned to batch_size. Commit: 96f78cd