[fx_importer] Convert non-persistent buffers lifted as tensor constants

python/torch_mlir/extras/fx_importer.py

@@ -342,13 +342,26 @@ def import_frozen_exported_program(self, prog: torch.export.ExportedProgram):
         sig = prog.graph_signature
         state_dict = prog.state_dict
         arg_replacements: dict[str, Any] = {}
-        # Lift buffers.
-        for input_name, state_name in sig.inputs_to_buffers.items():
-            try:
-                state_value = state_dict[state_name]
-            except KeyError as e:
-                raise AssertionError("Could not find state mapping for buffer") from e
-            arg_replacements[input_name] = state_value
+
+        # If there is no "constants" attribute, consult the "state_dict". Otherwise, only look
+        # at "constants". Relevant upstream patch: https://github.com/pytorch/pytorch/pull/118969
+        if hasattr(prog, "constants"):
+            constants = prog.constants
+            # Lift tensor constants.
+            for input_name, state_name in sig.inputs_to_lifted_tensor_constants.items():
+                try:
+                    state_value = constants[state_name]
+                except KeyError as e:
+                    raise AssertionError("Could not find state mapping for tensor constants") from e
+                arg_replacements[input_name] = state_value
+        else:
+            # Lift buffers.
+            for input_name, state_name in sig.inputs_to_buffers.items():
+                try:
+                    state_value = state_dict[state_name]
+                except KeyError as e:
+                    raise AssertionError("Could not find state mapping for buffer") from e
+                arg_replacements[input_name] = state_value
 
         # Lift parameters.
         for input_name, state_name in sig.inputs_to_parameters.items():

pytorch-hash.txt

@@ -1 +1 @@
-72fcb9ad662bb941a266e3d747835382634c2be6
+3cbc8e89fd09b0ffb4914187b438f15c121e2302

pytorch-requirements.txt

@@ -1,3 +1,3 @@
 -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
 --pre
-torch==2.3.0.dev20240122
+torch==2.3.0.dev20240207

torchvision-requirements.txt

@@ -1,3 +1,3 @@
 -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
 --pre
-torchvision==0.18.0.dev20240122
+torchvision==0.18.0.dev20240207