[torchlib] Implement torch.ops.prims.broadcast_in_dim.default

[Copilot]

This PR implements the missing torch.ops.prims.broadcast_in_dim.default operation that appears in BERT_pytorch and other PyTorch models.

Overview

The broadcast_in_dim operation is a primitive that broadcasts a tensor to a target shape by specifying which dimensions of the output correspond to the input tensor dimensions. This is different from standard broadcasting operations.

Implementation Details

Function signature:

def prims_broadcast_in_dim(
    a: TensorType, shape: INT64, broadcast_dimensions: Sequence[int]
) -> TensorType:

Parameters:

• a: Input tensor to broadcast
• shape: Target output shape
• broadcast_dimensions: Specifies which dimensions of the output shape correspond to the input tensor dimensions

Example:

# Input tensor: [3, 4]
# Target shape: [2, 3, 5, 4] 
# broadcast_dimensions: [1, 3]
# Result: Input dimension 0 (size 3) maps to output dimension 1
#         Input dimension 1 (size 4) maps to output dimension 3
#         Output dimensions 0 and 2 are broadcasted (filled from size 1)

Fixes #2218. Fix pytorch/pytorch#135343

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.33%. Comparing base (8ed3521) to head (690aa95).
⚠️ Report is 1 commits behind head on main.
✅ All tests successful. No failed tests found.

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2382      +/-   ##
==========================================
+ Coverage   70.31%   70.33%   +0.01%     
==========================================
  Files         222      222              
  Lines       26647    26657      +10     
  Branches     2664     2666       +2     
==========================================
+ Hits        18738    18749      +11     
+ Misses       6993     6992       -1     
  Partials      916      916              

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[github-advanced-security]

lintrunner found more than 20 potential problems in the proposed changes. Check the Files changed tab for more details.

[justinchuby]

What does broadcast_in_dim do?

[xadupre]

The proposed implementation is equivalent to the decomposition of broadcast_in_dims in pytorch.

[justinchuby]

Nice: I didn't know it exsited: https://github.com/pytorch/pytorch/blob/799471d92b15c433a0a92e5c4cfb11350841b123/torch/_inductor/lowering.py#L1111-L1121

[justinchuby]

Will find ways to test as a follow up

[titaiwangms]

Is there a test case?

[justinchuby]

Added tests