@@ -1003,37 +1003,6 @@ def forward(self, mat1: Tensor, mat2: Tensor) -> Tensor:
10031003 model = model , mat1 = mat1 , mat2 = mat2 , dynamic_shapes = dynamic_shapes
10041004 )
10051005
1006- async def test_mixed_dtypes (self ) -> None :
1007- """Test bmm with mixed f32/f16 inputs.
1008-
1009- Reproduces the EfficientSam pattern: model.half() makes weights f16,
1010- but an explicit dtype=torch.float32 tensor creates f32 that flows
1011- into a bmm with f16 weights.
1012- """
1013-
1014- class MixedBmmModel (nn .Module ):
1015- def __init__ (self ) -> None :
1016- super ().__init__ ()
1017- self .weight = nn .Parameter (torch .randn (3 , 8 , 4 ))
1018-
1019- def forward (self , x : Tensor ) -> Tensor :
1020- # Explicit f32 creation contaminates x via add
1021- f32_val = torch .ones (1 , device = x .device , dtype = torch .float32 )
1022- x = x + f32_val # promotes x(f16) to f32
1023- return torch .bmm (x , self .weight ) # f32 @ f16
1024-
1025- model = MixedBmmModel ().eval ().half ()
1026- x = torch .randn (3 , 4 , 8 , dtype = torch .float16 )
1027-
1028- with torch .autocast (device_type = "cpu" , dtype = torch .float16 ):
1029- exported_program = torch .export .export (model , args = (), kwargs = {"x" : x })
1030- exported_program = exported_program .run_decompositions (
1031- torch .export .default_decompositions ()
1032- )
1033-
1034- converter = TorchConverter ().add_exported_program (exported_program )
1035- converter .to_coreai ()
1036-
10371006
10381007class TestCat :
10391008 """Test suite for aten.cat → coreai.concat conversion."""
0 commit comments