Skip to content

Add new_kwargs function to fix dtype generalization pass #614

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Xreki merged 2 commits into from
Jan 28, 2026
+22 −8
Merged

Add new_kwargs function to fix dtype generalization pass #614

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7b43173
Add FP32_ONLY_FUNCS op to fix dtype generalization pass
WHoutstanding Jan 27, 2026
e269e2e
fix bug of invalid dtype for bias
WHoutstanding Jan 28, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 22 additions & 8 deletions graph_net/torch/dtype_gen_passes/dtype_generalization_pass.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -94,6 +94,10 @@ def create_placeholder(node: fx.Node) -> fx.Node:
"""Create a placeholder node with dtype conversion if needed."""
new_node = new_graph.node_copy(node, lambda x: val_map.get(x, x))
if self._is_float32_tensor(node):
attr_name = str(node.target)
if self.should_preserve_weight(attr_name):
return new_node

return new_graph.call_method("to", args=(new_node, self.torch_dtype))
return new_node

Expand Down Expand Up @@ -121,17 +125,29 @@ def create_new_args(node: fx.Node) -> list:
new_args.append(arg)
return new_args

def create_new_kwargs(node: fx.Node) -> dict:
"""new_kwargs of node with dtype conversion if needed."""
new_kwargs = {}

for k, v in node.kwargs.items():
if isinstance(v, fx.Node):
mapped = val_map[v]
if self._is_float32_tensor(v):
mapped = new_graph.call_method("to", (mapped, self.torch_dtype))
else:
new_kwargs[k] = mapped
else:
new_kwargs[k] = v
return new_kwargs

def create_call_function(node: fx.Node) -> fx.Node:
"""Create a call_function node with dtype conversion if needed."""
if node.target not in AMP_CALL_FUNCTION:
return new_graph.node_copy(node, lambda x: val_map[x])

new_args = create_new_args(node)

new_kwargs = {
k: val_map[v] if isinstance(v, fx.Node) else v
for k, v in node.kwargs.items()
}
new_kwargs = create_new_kwargs(node)

return new_graph.call_function(
node.target,
Expand All @@ -140,15 +156,13 @@ def create_call_function(node: fx.Node) -> fx.Node:
)

def create_call_method(node: fx.Node) -> fx.Node:
"""Create a call_method node with dtype conversion if needed."""
if node.target not in AMP_CALL_METHOD:
return new_graph.node_copy(node, lambda x: val_map[x])

new_args = create_new_args(node)

new_kwargs = {
k: (val_map[v] if isinstance(v, fx.Node) else v)
for k, v in node.kwargs.items()
}
new_kwargs = create_new_kwargs(node)

return new_graph.call_method(
node.target,
Expand Down