llvm
/
torch-mlir
mirror of https://github.com/llvm/torch-mlir
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

[FxImporter] Add InputInfo to Resolve Literal Hook (#3688)

pull/3689/head
Christopher McGirr 2024-09-05 18:53:11 +02:00 committed by GitHub
parent 295bf418a4
commit b790061b69
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 4 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
python/torch_mlir/extras/fx_importer.py
View File

@ -470,7 +470,7 @@ class FxImporterHooks:
... ...
def resolve_literal( def resolve_literal(
self, gni: "GraphNodeImporter", literal: Any self, gni: "GraphNodeImporter", literal: Any, info: Optional[InputInfo]
) -> Optional[Value]: ) -> Optional[Value]:
"""User overridable hook to resolve a literal value.""" """User overridable hook to resolve a literal value."""
return None return None
@ -1826,13 +1826,13 @@ class GraphNodeImporter:
name=op_name, results=[result_type], operands=operands name=op_name, results=[result_type], operands=operands
).result ).result
def _import_literal(self, py_value: Any) -> Value: def _import_literal(self, py_value: Any, info: Optional[InputInfo] = None) -> Value:
orig_value = None orig_value = None
if isinstance(py_value, torch.Tensor) and py_value.dtype == torch.bool: if isinstance(py_value, torch.Tensor) and py_value.dtype == torch.bool:
orig_value = py_value orig_value = py_value
py_value = py_value.to(torch.uint8) py_value = py_value.to(torch.uint8)
# Apply the conversion callback. # Apply the conversion callback.
user_value = self.fx_importer._hooks.resolve_literal(self, py_value) user_value = self.fx_importer._hooks.resolve_literal(self, py_value, info)
if user_value is not None: if user_value is not None:
assert isinstance(user_value, Value) assert isinstance(user_value, Value)
if orig_value is not None: if orig_value is not None:
@ -1866,7 +1866,7 @@ class GraphNodeImporter:
raise ValueError( raise ValueError(
f"Cannot import {info.input_spec} as a literal because it is mutable" f"Cannot import {info.input_spec} as a literal because it is mutable"
) )
return self._import_literal(py_value) return self._import_literal(py_value, info)
def _import_scalar_as_tensor(self, loc: Location, arg: NodeArgument) -> Value: def _import_scalar_as_tensor(self, loc: Location, arg: NodeArgument) -> Value:
tensor_arg = torch.tensor(arg) tensor_arg = torch.tensor(arg)