diff --git a/py/torch_tensorrt/dynamo/_exporter.py b/py/torch_tensorrt/dynamo/_exporter.py index c6c992d720..74c819fe73 100644 --- a/py/torch_tensorrt/dynamo/_exporter.py +++ b/py/torch_tensorrt/dynamo/_exporter.py @@ -888,6 +888,20 @@ def inline_trt_modules( """ Replace TRT submodules with trt engine nodes. """ + fake_mode = detect_fake_mode( + tuple( + node.meta["val"] + for node in gm.graph.nodes + if node.op == "placeholder" and "val" in node.meta + ) + ) + + def scalar_tensor_meta(dtype: torch.dtype, device: torch.device) -> torch.Tensor: + if fake_mode is None: + return torch.empty((), dtype=dtype, device="meta") + with fake_mode: + return torch.empty((), dtype=dtype, device=device) + for name, _ in gm.named_children(): if "_run_on_acc" not in name: continue @@ -904,16 +918,52 @@ def inline_trt_modules( raise ValueError( f"trt_module_node: {trt_module_node.name} does not have the metadata which should be set during dynamo compile_module step." ) - num_outputs = len(trt_module_node.meta["val"]) + + shape_info = trt_module.symbolic_shape_expressions or {} + input_info = shape_info.get("inputs", []) + output_info = shape_info.get("outputs", []) + original_output_vals = trt_module_node.meta["val"] + if not isinstance(original_output_vals, (tuple, list)): + original_output_vals = [original_output_vals] + num_outputs = len(original_output_vals) + + engine_output_vals = [] + for index, value in enumerate(original_output_vals): + if index < len(output_info) and output_info[index].get("is_scalar"): + engine_output_vals.append( + scalar_tensor_meta( + output_info[index]["dtype"], trt_module.target_device + ) + ) + else: + engine_output_vals.append(value) + # Insert a call_function node to perform inference on TRT engine with gm.graph.inserting_before(trt_module_node): + engine_inputs = list(trt_module_node.args) + for index, info in enumerate(input_info): + if index >= len(engine_inputs) or not info.get("is_scalar"): + continue + scalar_input = gm.graph.call_function( + torch.ops.aten.scalar_tensor.default, + (engine_inputs[index],), + { + "dtype": info["dtype"], + "device": trt_module.target_device, + }, + ) + scalar_input.meta["val"] = scalar_tensor_meta( + info["dtype"], trt_module.target_device + ) + engine_inputs[index] = scalar_input + if cross_compile_module: engine_info = trt_module._pack_engine_info() engine_bytes = engine_info[ENGINE_IDX] engine_info[ENGINE_IDX] = base64.b64encode(engine_bytes).decode("utf-8") trt_node = gm.graph.call_function( torch.ops.tensorrt.no_op_placeholder_for_execute_engine.default, - (trt_module_node.args, *engine_info), + (tuple(engine_inputs), *engine_info), ) else: engine_name = f"{name}_engine" @@ -922,28 +972,50 @@ def inline_trt_modules( trt_node = gm.graph.call_function( torch.ops.tensorrt.execute_engine.default, - (trt_module_node.args, engine_node), + (tuple(engine_inputs), engine_node), ) engine_node.meta["val"] = CustomObjArgument( name=engine_node.name, class_fqn="" ) assert num_outputs > 0 - trt_node.meta["val"] = trt_module_node.meta["val"] + trt_node.meta["val"] = engine_output_vals + + def restore_scalar_output( + getitem_node: torch.fx.Node, output_index: int + ) -> torch.fx.Node: + getitem_node.meta["val"] = engine_output_vals[output_index] + if output_index >= len(output_info) or not output_info[output_index].get( + "is_scalar" + ): + return getitem_node + + with gm.graph.inserting_after(getitem_node): + scalar_output = gm.graph.call_function( + torch.ops.aten.item.default, (getitem_node,) + ) + scalar_output.meta["val"] = original_output_vals[output_index] + getitem_node.replace_all_uses_with( + scalar_output, delete_user_cb=lambda user: user is not scalar_output + ) + return scalar_output if num_outputs == 1: - # Insert getitem nodes as outputs (for export serialization to work) + # Insert a getitem because execute_engine always returns Tensor[]. with gm.graph.inserting_after(trt_node): getitem_output = gm.graph.call_function(operator.getitem, (trt_node, 0)) - getitem_output.meta["val"] = trt_node.meta["val"] - trt_module_node.replace_all_uses_with(getitem_output) + replacement = restore_scalar_output(getitem_output, 0) + trt_module_node.replace_all_uses_with(replacement) else: - # Multiple outputs case: - # Replace uses of submodule with the trt_node. - # getitem nodes are already added inherently by the partitioner + # Multiple-output partitioner graphs already contain getitem users. trt_module_node.replace_all_uses_with(trt_node) - getitem_nodes = trt_node.users - for idx, getitem_node in enumerate(getitem_nodes): - getitem_node.meta["val"] = trt_node.meta["val"][idx] + for getitem_node in list(trt_node.users): + if ( + getitem_node.target is not operator.getitem + or len(getitem_node.args) < 2 + or not isinstance(getitem_node.args[1], int) + ): + continue + restore_scalar_output(getitem_node, getitem_node.args[1]) # Expose the engine's aliased (KV-cache) outputs as graph-level buffer # mutations so the ExecuTorch path sees a real mutable buffer instead of diff --git a/py/torch_tensorrt/dynamo/conversion/_conversion.py b/py/torch_tensorrt/dynamo/conversion/_conversion.py index 38d104473a..f70b732cbd 100644 --- a/py/torch_tensorrt/dynamo/conversion/_conversion.py +++ b/py/torch_tensorrt/dynamo/conversion/_conversion.py @@ -4,6 +4,7 @@ import logging from typing import Any, Dict, List, NamedTuple, Optional, Sequence, Tuple +import tensorrt as trt import torch from torch_tensorrt._enums import dtype from torch_tensorrt._features import ENABLED_FEATURES @@ -25,8 +26,6 @@ ) from torch_tensorrt.logging import TRT_LOGGER -import tensorrt as trt - logger = logging.getLogger(__name__) @@ -221,7 +220,9 @@ def interpret_module_to_result( SerializedInterpreterResult """ - symbolic_shape_expressions = extract_symbolic_shape_expressions(module) + symbolic_shape_expressions = extract_symbolic_shape_expressions( + module, inputs=inputs + ) if symbolic_shape_expressions is None: raise RuntimeError( "Failed to extract symbolic shape expressions from source FX graph partition" diff --git a/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py b/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py index e8f84921a2..11213d9d98 100644 --- a/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py +++ b/py/torch_tensorrt/dynamo/conversion/_symbolic_shape_capture.py @@ -7,15 +7,18 @@ """ import logging -from typing import Any, Dict, List, Optional +from typing import Any, Dict, List, Optional, Sequence import torch +from torch_tensorrt._enums import dtype as _dtype +from torch_tensorrt._Input import Input logger = logging.getLogger(__name__) def extract_symbolic_shape_expressions( module: torch.fx.GraphModule, + inputs: Optional[Sequence[Input]] = None, ) -> Optional[Dict[str, List[Dict[str, Any]]]]: """ Extract symbolic shape expressions from an FX graph. @@ -25,11 +28,23 @@ def extract_symbolic_shape_expressions( Args: module: FX GraphModule with symbolic shapes in node metadata + inputs: Engine Input specs from construct_submodule_inputs, aligned by + name. Used as the dtype source of truth for scalar inputs, which + have no dtype of their own in FX metadata. Falls back to a + best-effort default when not provided. Returns: Dict with 'inputs' and 'outputs' keys, each containing a list of dicts with shape_exprs and dtype, or None if extraction fails """ + # dtype.unknown has no torch.dtype equivalent; skip it rather than raise + # (unset-dtype inputs are never looked up below anyway). + input_dtypes_by_name = { + inp.name: inp.dtype.to(torch.dtype) + for inp in inputs or () + if inp.dtype != _dtype.unknown + } + # Find input nodes (placeholders) input_nodes = [node for node in module.graph.nodes if node.op == "placeholder"] @@ -70,11 +85,15 @@ def extract_symbolic_shape_expressions( ) elif isinstance(input_val, (torch.SymInt, torch.SymFloat, int, float, bool)): if isinstance(input_val, (torch.SymInt, int)): - scalar_dtype = torch.int64 + default_scalar_dtype = torch.int64 elif isinstance(input_val, (torch.SymFloat, float)): - scalar_dtype = torch.float64 + default_scalar_dtype = torch.float32 else: - scalar_dtype = torch.bool + default_scalar_dtype = torch.bool + # Prefer the engine's actual binding dtype over the guess above. + scalar_dtype = input_dtypes_by_name.get( + input_node.name, default_scalar_dtype + ) input_info.append( { "shape_exprs": [], @@ -122,7 +141,8 @@ def extract_symbolic_shape_expressions( if isinstance(out_val, (torch.SymInt, int)): scalar_dtype = torch.int64 elif isinstance(out_val, (torch.SymFloat, float)): - scalar_dtype = torch.float64 + # No float64 output binding exists in TensorRT. + scalar_dtype = torch.float32 else: scalar_dtype = torch.bool output_info.append( diff --git a/py/torch_tensorrt/dynamo/partitioning/common.py b/py/torch_tensorrt/dynamo/partitioning/common.py index 62e68449b6..a0942a5d0b 100644 --- a/py/torch_tensorrt/dynamo/partitioning/common.py +++ b/py/torch_tensorrt/dynamo/partitioning/common.py @@ -283,9 +283,11 @@ def construct_submodule_inputs( ) ) elif isinstance(input_meta, torch.SymFloat): + # Rank-0 to match the 0-D scalar_tensor inline_trt_modules + # materializes for this input at the engine boundary. torchtrt_inputs.append( get_input( - [1], + [], torch.float32, name=input.name, is_shape_tensor=False, # Only SymInt inputs are treated as shape tensors diff --git a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py index b52bb8d360..6c2f8c7e9e 100644 --- a/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py +++ b/py/torch_tensorrt/dynamo/runtime/_TorchTensorRTModule.py @@ -711,7 +711,7 @@ def pre_allocated_outputs(self) -> Any: def set_use_output_allocator(self, enable: bool) -> None: self.get_engine().use_output_allocator_outputs = enable - def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]: + def forward(self, *inputs: Any) -> Any: """Run the TensorRT engine on GPU tensors (non-tensor args are cast to CUDA tensors). Note: callers are responsible for ensuring the engine has been set up; @@ -745,7 +745,7 @@ def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]: else: input_tensors.append(torch.tensor(i).cuda()) - outputs: List[torch.Tensor] = torch.ops.tensorrt.execute_engine( + outputs: List[Any] = torch.ops.tensorrt.execute_engine( list(input_tensors), self.engine ) @@ -760,6 +760,11 @@ def forward(self, *inputs: Any) -> torch.Tensor | Tuple[torch.Tensor, ...]: if n < len(outputs): outputs = outputs[:n] + output_info = (self.symbolic_shape_expressions or {}).get("outputs", []) + for index, info in enumerate(output_info): + if index < len(outputs) and info.get("is_scalar"): + outputs[index] = outputs[index].item() + if len(outputs) == 1: return outputs[0] diff --git a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py index fcb6f21eb8..294e8ab72a 100644 --- a/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py +++ b/py/torch_tensorrt/dynamo/runtime/meta_ops/register_meta_ops.py @@ -10,7 +10,7 @@ def _apply_symbolic_shape_expressions( - inputs: List[torch.Tensor], shape_info: Dict[str, List[Dict[str, Any]]] + inputs: List[Any], shape_info: Dict[str, List[Dict[str, Any]]] ) -> List[torch.Tensor]: """ Apply symbolic shape expressions to create output fake tensors. @@ -35,6 +35,13 @@ def _apply_symbolic_shape_expressions( input_info = shape_info.get("inputs", []) output_info = shape_info.get("outputs", []) + tensor_inputs = [value for value in inputs if isinstance(value, torch.Tensor)] + if not tensor_inputs: + raise RuntimeError( + "[torch.ops.tensorrt.execute_engine]: At least one tensor input is required to infer the engine output device" + ) + output_device = tensor_inputs[0].device + fake_mode = detect_fake_mode(inputs) if fake_mode is None: # No fake mode - shouldn't happen, but fall back to concrete shapes @@ -45,7 +52,7 @@ def _apply_symbolic_shape_expressions( for s in info["shape_exprs"] ] outputs.append( - torch.empty(shape, dtype=info["dtype"], device=inputs[0].device) + torch.empty(shape, dtype=info["dtype"], device=output_device) ) return outputs @@ -82,6 +89,8 @@ def in_compile_namespace(expr: sympy.Expr) -> sympy.Expr: # Align inputs: for each captured input, match it with the corresponding runtime input for inp_tensor, inp_info in zip(inputs, input_info): + if inp_info.get("is_scalar"): + continue for d, compile_expr in zip(inp_tensor.shape, inp_info["shape_exprs"]): if isinstance(compile_expr, int): continue @@ -226,7 +235,7 @@ def in_compile_namespace(expr: sympy.Expr) -> sympy.Expr: ) from e outputs.append( - torch.empty(output_shape, dtype=info["dtype"], device=inputs[0].device) + torch.empty(output_shape, dtype=info["dtype"], device=output_device) ) logger.debug( f"[torch.ops.tensorrt.execute_engine]: Meta kernel found the following output FakeTensors: {outputs}" diff --git a/tests/py/dynamo/conversion/test_symbolic_shape_capture.py b/tests/py/dynamo/conversion/test_symbolic_shape_capture.py new file mode 100644 index 0000000000..8a84661117 --- /dev/null +++ b/tests/py/dynamo/conversion/test_symbolic_shape_capture.py @@ -0,0 +1,76 @@ +import torch +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv +from torch.testing._internal.common_utils import TestCase, run_tests +from torch_tensorrt import Input +from torch_tensorrt.dynamo.conversion._symbolic_shape_capture import ( + extract_symbolic_shape_expressions, +) + + +class TestSymbolicShapeCaptureScalarDtype(TestCase): + """A SymFloat scalar input has no dtype of its own in FX metadata; the + recorded dtype must match the engine's actual binding (float32).""" + + def _make_symfloat_input_module(self) -> torch.fx.GraphModule: + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + scale = shape_env.create_unbacked_symfloat() + x = torch.empty((4,), dtype=torch.float32) + + graph = torch.fx.Graph() + scale_input = graph.placeholder("scale") + scale_input.meta["val"] = scale + x_input = graph.placeholder("x") + x_input.meta["val"] = x + output = graph.call_function( + torch.ops.aten.mul.Tensor, args=(x_input, scale_input) + ) + output.meta["val"] = x + graph.output(output) + return torch.fx.GraphModule({}, graph) + + def test_symfloat_input_uses_engine_binding_dtype(self): + module = self._make_symfloat_input_module() + engine_inputs = [ + Input([1], dtype=torch.float32, name="scale"), + Input((4,), dtype=torch.float32, name="x"), + ] + + metadata = extract_symbolic_shape_expressions(module, inputs=engine_inputs) + + scale_info = next( + info for info in metadata["inputs"] if info["name"] == "scale" + ) + self.assertTrue(scale_info["is_scalar"]) + self.assertEqual(scale_info["dtype"], torch.float32) + + def test_symfloat_input_defaults_to_float32_without_engine_inputs(self): + # Used to hardcode float64 -- no real engine binding is ever float64. + module = self._make_symfloat_input_module() + + metadata = extract_symbolic_shape_expressions(module) + + scale_info = next( + info for info in metadata["inputs"] if info["name"] == "scale" + ) + self.assertEqual(scale_info["dtype"], torch.float32) + + def test_ordinary_input_without_explicit_dtype_does_not_raise(self): + # An Input with dtype omitted (dtype.unknown) used to raise TypeError. + module = self._make_symfloat_input_module() + engine_inputs = [ + Input([1], dtype=torch.float32, name="scale"), + Input((4,)), # dtype intentionally omitted + ] + + metadata = extract_symbolic_shape_expressions(module, inputs=engine_inputs) + + scale_info = next( + info for info in metadata["inputs"] if info["name"] == "scale" + ) + self.assertEqual(scale_info["dtype"], torch.float32) + + +if __name__ == "__main__": + run_tests() diff --git a/tests/py/dynamo/models/test_exporter_inlining.py b/tests/py/dynamo/models/test_exporter_inlining.py index 2fa2f0470e..2029662569 100644 --- a/tests/py/dynamo/models/test_exporter_inlining.py +++ b/tests/py/dynamo/models/test_exporter_inlining.py @@ -6,6 +6,8 @@ import pytest import torch +from torch._subclasses.fake_tensor import FakeTensorMode +from torch.fx.experimental.symbolic_shapes import ShapeEnv from torch.export.graph_signature import ( ExportGraphSignature, InputKind, @@ -17,6 +19,7 @@ from torch_tensorrt.dynamo._exporter import ( create_trt_exp_program, inline_torch_modules, + inline_trt_modules, lift, ) @@ -147,6 +150,96 @@ def test_inline_torch_modules_computed_intermediate_inputs(): assert torch.allclose(out, torch.tensor(16.0)) +class _FakeTRTModule(torch.nn.Module): + def __init__(self, shape_info): + super().__init__() + self.engine = object() + self.symbolic_shape_expressions = shape_info + self.target_device = torch.device("cuda") + self.aliased_io = {} + + def forward(self, *args): + raise AssertionError("synthetic TRT module should be inlined, not executed") + + +@pytest.mark.unit +def test_inline_trt_modules_tensorizes_scalar_inputs_and_restores_outputs(): + """The inlined execute_engine Tensor[] boundary contains tensors only. + + Scalar partition inputs are materialized as zero-dimensional tensors before + execute_engine, and outputs recorded as scalars are converted back with item + before reaching their original consumers. + """ + shape_info = { + "inputs": [ + { + "shape_exprs": [], + "dtype": torch.int64, + "name": "length", + "is_scalar": True, + }, + { + "shape_exprs": [4], + "dtype": torch.float32, + "name": "x", + }, + ], + "outputs": [ + { + "shape_exprs": [], + "dtype": torch.int64, + "is_scalar": True, + }, + { + "shape_exprs": [4], + "dtype": torch.float32, + }, + ], + } + + shape_env = ShapeEnv() + fake_mode = FakeTensorMode(shape_env=shape_env) + with fake_mode: + scalar_val = shape_env.create_unbacked_symint() + tensor_val = torch.empty(4, device="cuda") + + graph = torch.fx.Graph() + scalar = graph.placeholder("length") + scalar.meta["val"] = scalar_val + tensor = graph.placeholder("x") + tensor.meta["val"] = tensor_val + + root = torch.nn.Module() + root.add_module("_run_on_acc_0", _FakeTRTModule(shape_info)) + call = graph.call_module("_run_on_acc_0", (scalar, tensor)) + call.meta["val"] = [scalar_val, tensor_val] + scalar_output = graph.call_function(operator.getitem, (call, 0)) + scalar_output.meta["val"] = scalar_val + tensor_output = graph.call_function(operator.getitem, (call, 1)) + tensor_output.meta["val"] = tensor_val + graph.output((scalar_output, tensor_output)) + gm = torch.fx.GraphModule(root, graph) + + inline_trt_modules(gm, expose_aliased_mutations=False) + gm.graph.lint() + + execute = next( + node + for node in gm.graph.nodes + if node.target is torch.ops.tensorrt.execute_engine.default + ) + engine_inputs = execute.args[0] + assert engine_inputs[0].target is torch.ops.aten.scalar_tensor.default + assert all(isinstance(node.meta.get("val"), torch.Tensor) for node in engine_inputs) + assert all(isinstance(value, torch.Tensor) for value in execute.meta["val"]) + + item_nodes = [ + node for node in gm.graph.nodes if node.target is torch.ops.aten.item.default + ] + assert len(item_nodes) == 1 + assert item_nodes[0].meta["val"] is scalar_val + + @pytest.mark.unit def test_create_trt_exp_program_rebuilds_in_spec_without_inputs(): """create_trt_exp_program must rebuild a correct in_spec on the plain-CodeGen diff --git a/tests/py/dynamo/models/test_meta_kernel_shape_inference.py b/tests/py/dynamo/models/test_meta_kernel_shape_inference.py index a0d6f731fd..4230c9b0ac 100644 --- a/tests/py/dynamo/models/test_meta_kernel_shape_inference.py +++ b/tests/py/dynamo/models/test_meta_kernel_shape_inference.py @@ -484,6 +484,58 @@ def test_rejects_inconsistent_repeated_direct_mapping(self): with pytest.raises(RuntimeError): _apply_symbolic_shape_expressions([fake_x, fake_z], shape_info) + @staticmethod + def _scalar_first_shape_info(): + return { + "inputs": [ + { + "shape_exprs": [], + "dtype": torch.int64, + "name": "length", + "is_scalar": True, + }, + { + "shape_exprs": [4], + "dtype": torch.float32, + "name": "x", + }, + ], + "outputs": [ + { + "shape_exprs": [4], + "dtype": torch.float32, + } + ], + } + + def test_scalar_input_is_not_walked_as_a_tensor(self): + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + scalar = shape_env.create_unbacked_symint() + tensor = torch.empty(4, device="cuda") + output = _apply_symbolic_shape_expressions( + [scalar, tensor], self._scalar_first_shape_info() + )[0] + + assert output.shape == (4,) + + def test_output_device_comes_from_first_tensor_input(self): + shape_env = ShapeEnv() + with FakeTensorMode(shape_env=shape_env): + scalar = shape_env.create_unbacked_symint() + tensor = torch.empty(4, device="cuda") + fake_output = _apply_symbolic_shape_expressions( + [scalar, tensor], self._scalar_first_shape_info() + )[0] + + concrete_output = _apply_symbolic_shape_expressions( + [shape_env.create_unbacked_symint(), torch.empty(4)], + self._scalar_first_shape_info(), + )[0] + + assert fake_output.device == tensor.device + assert concrete_output.device.type == "cpu" + if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) diff --git a/tests/py/dynamo/models/test_symint_scalar_input.py b/tests/py/dynamo/models/test_symint_scalar_input.py index 9ab3104064..66f366c05c 100644 --- a/tests/py/dynamo/models/test_symint_scalar_input.py +++ b/tests/py/dynamo/models/test_symint_scalar_input.py @@ -14,6 +14,7 @@ import pytest import torch import torch_tensorrt as torchtrt +from torch_tensorrt.dynamo._exporter import transform from torch_tensorrt.dynamo.utils import COSINE_THRESHOLD, cosine_similarity assertions = unittest.TestCase() @@ -194,3 +195,55 @@ def forward(self, x, targets): ) torch._dynamo.reset() + + +@pytest.mark.unit +def test_symfloat_scalar_input(): + """ + A data-dependent Python float (e.g. weights.sum().item()) crossing a TRT + partition boundary as a SymFloat must be materialized with the engine's + actual binding dtype (float32), not a guessed one. + """ + torch._dynamo.config.capture_scalar_outputs = True + try: + + class ScaleByDataDependentFloat(torch.nn.Module): + def forward(self, x, weights): + scale = weights.sum().item() + return x * scale + + model = ScaleByDataDependentFloat().eval().cuda() + x = torch.randn(8).cuda() + weights = torch.randn(4).cuda() + expected = model(x, weights) + + exported = torch.export.export(model, (x, weights)) + compiled = torchtrt.dynamo.compile( + exported, + inputs=(x, weights), + min_block_size=1, + pass_through_build_failures=True, + ) + + # The inserted scalar tensor must carry the engine binding's dtype. + inlined = transform(compiled) + inlined.recompile() + scalar_tensor_nodes = [ + node + for node in inlined.graph.nodes + if node.target == torch.ops.aten.scalar_tensor.default + ] + assertions.assertEqual(len(scalar_tensor_nodes), 1) + assertions.assertEqual( + scalar_tensor_nodes[0].kwargs["dtype"], + torch.float32, + msg="scalar_tensor inserted for the SymFloat input must match the " + "float32 engine binding dtype, not the Python-type-based default", + ) + + # And the exported module must actually execute successfully. + actual = compiled(x, weights) + torch.testing.assert_close(actual, expected) + finally: + torch._dynamo.config.capture_scalar_outputs = False + torch._dynamo.reset()