feat: modify aten_bilinear from einsum to matmul

[fw7th]

What this does

• Replaces Bilinear Transformation's einsum implementation with matmul

Why

• Performance increase from MatMul over Einsum

Testing

• Existing tests pass for function implementation

Resolves #2573

[fw7th]

microsoft

@microsoft-github-policy-service agree

[fw7th]

Note on test failure

The failure appears unrelated to the bilinear changes. The error is raised during test input generation for upsample_bilinear2d_* in full-graph mode
(ValueError: Seed must be between 0 and 2**32 - 1).

This suggests an RNG seed issue in the test harness rather than an op
implementation problem. I can reproduce it consistently locally.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 70.11%. Comparing base (70abdd9) to head (c7e30e6).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2746   +/-   ##
=======================================
  Coverage   70.10%   70.11%           
=======================================
  Files         228      228           
  Lines       27394    27403    +9     
  Branches     2781     2781           
=======================================
+ Hits        19204    19213    +9     
  Misses       7234     7234           
  Partials      956      956           

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

[justinchuby]

Thanks. @xadupre or @gramalingam for another review

[Copilot]

Pull request overview

This PR aims to optimize the aten_bilinear function by replacing the Einstein summation (einsum) implementation with matrix multiplication (matmul) operations for improved performance.

Key changes:

• Replaced single einsum operation with a sequence of transpose, reshape, matmul, and squeeze operations
• Added explicit shape extraction for batch dimensions and feature dimensions
• Modified computation flow to use matmul instead of einsum for the bilinear transformation

[Copilot]

The variable names in the comments don't match the actual variable names used in the code. The comment mentions 'B' for batch dimension, but the code uses 'batch_size'. Similarly, the comment uses 'H1', 'H2', and 'O' which aren't actual variables. For clarity, either use the actual variable names in comments or define these as actual variables in the code.

Suggested change

	# (O, H1, H2) -> (H1, O, H2) where O = output_shape, H1 = input1_shape, H2 = input2_shape
	W_permute = op.Transpose(weight, perm=[1, 0, 2])
	# (H1, O, H2) -> (H1, O * H2)
	W_flat = op.Reshape(
	W_permute,
	op.Concat(input1_shape, op.Mul(output_shape, input2_shape), axis=0),
	)
	# (B, H1) @ (H1, O*H2) -> (B, O*H2)
	tmp = op.MatMul(input1, W_flat)
	# (B, O*H2) -> (B, O, H2)
	tmp = op.Reshape(tmp, op.Concat(batch_size, output_shape, input2_shape, axis=0))
	# (B, H2) -> (B, H2, 1) -> (B, O, H2) @ (B, H2, 1) -> (B, O, 1) -> (B, O)
	# (out_features, in1_features, in2_features) -> (in1_features, out_features, in2_features)
	W_permute = op.Transpose(weight, perm=[1, 0, 2])
	# (in1_features, out_features, in2_features) -> (in1_features, out_features * in2_features)
	W_flat = op.Reshape(
	W_permute,
	op.Concat(input1_shape, op.Mul(output_shape, input2_shape), axis=0),
	)
	# (batch_size, in1_features) @ (in1_features, out_features * in2_features)
	# -> (batch_size, out_features * in2_features)
	tmp = op.MatMul(input1, W_flat)
	# (batch_size, out_features * in2_features)
	# -> (batch_size, out_features, in2_features)
	tmp = op.Reshape(tmp, op.Concat(batch_size, output_shape, input2_shape, axis=0))
	# (batch_size, in2_features) -> (batch_size, in2_features, 1)
	# -> (batch_size, out_features, in2_features) @ (batch_size, in2_features, 1)
	# -> (batch_size, out_features, 1) -> (batch_size, out_features)

[Copilot]

The test cases in sample_inputs_bilinear only test 2D inputs with shape (batch_size, features), but the bilinear function is documented to support arbitrary batch dimensions with shape (..., features). The test suite should include cases with multi-dimensional batch inputs like (B1, B2, in_features) to ensure the implementation correctly handles these cases. This would have caught the current implementation bug where only the first batch dimension is extracted.

[Copilot]

This line only extracts the first dimension of input1, but the bilinear function should support arbitrary batch dimensions (as indicated by the "..." notation in the comments above). The original einsum implementation correctly handled multiple batch dimensions. To fix this, you should extract all batch dimensions except the last one, which would be op.Shape(input1, end=-1).

Suggested change

	batch_size = op.Shape(input1, start=0, end=1)
	batch_size = op.Shape(input1, end=-1)

[Copilot]

This MatMul operation assumes input1 has shape (B, H1) with only a single batch dimension, but input1 can have arbitrary batch dimensions like (B1, B2, ..., Bn, H1). The matmul needs to handle multi-dimensional batch inputs correctly. When input1 has shape (..., H1) where ... represents multiple dimensions, this needs to be flattened to (*, H1) before the matmul, or the weight needs to be broadcasted appropriately.

[Copilot]

This reshape only uses batch_size (the first dimension), but should use all batch dimensions. For inputs with shape (..., H1), the result of the previous matmul would have shape (, OH2) where * represents the flattened batch dimensions, and this needs to be reshaped to (..., O, H2) where ... represents the original multi-dimensional batch shape.

[Copilot]

This operation assumes input2 has shape (B, H2) with a single batch dimension, but input2 can have arbitrary batch dimensions (..., H2). The Unsqueeze and subsequent MatMul operations will fail or produce incorrect results when input2 has multi-dimensional batch shape. For example, if input2 has shape (B1, B2, H2), the unsqueeze would create (B1, B2, H2, 1), but tmp has shape (B1, O, H2), making the matmul dimension-incompatible.

[justinchuby]

Output has shape (..., out_features). The final output shape should be op.Concat(op.Shape(input1, start=0, end=-1), op.Shape(weight, start=0, end=1))

[justinchuby]

Instead of using batch_size, I think you should use op.Shape(input1, start=0, end=-1) to get all of the leading shapes.

[justinchuby]

Found some issues