Skip to content

Commit 47bf2ea

Browse files
authored
Promote in convolution (#5727)
Currently this model: ``` m = timm.create_model("efficientformerv2_s0", pretrained=True, scriptable=True).eval()` ``` will fail to convert to stablehlo because we use mixed shape in Conv. this is allowed in HLO but not in mhlo. Workaroudn by manually promoting
1 parent 294610a commit 47bf2ea

File tree

1 file changed

+2
-1
lines changed

1 file changed

+2
-1
lines changed

torch_xla/csrc/convolution.cpp

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -348,7 +348,8 @@ xla::XlaOp BuildConvolutionOverrideableBias(
348348
xla::XlaOp bias_broadcast =
349349
xla::Transpose(xla::Broadcast(bias, broadcast_sizes),
350350
BiasTransposePermutation(broadcast_sizes.size() + 1));
351-
return conv + bias_broadcast;
351+
auto promoted = XlaHelpers::Promote(conv, bias_broadcast);
352+
return promoted.first + promoted.second;
352353
}
353354

354355
ConvGrads BuildConvolutionBackwardOverrideable(

0 commit comments

Comments
 (0)