Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
21 changes: 21 additions & 0 deletions graph_net/torch/backend/unstable_to_stable_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -126,6 +126,27 @@ def _impl_unstable_to_stable_fftn(self, gm):

return gm

def _impl_unstable_to_stable_special_logit(self, gm):
"""
Convert torch._C._special.special_logit to torch.special.logit
"""
issue_nodes = (
node
for node in gm.graph.nodes
if node.op == "call_function"
if hasattr(node.target, "__module__")
if node.target.__module__ == "torch._C._special"
if hasattr(node.target, "__name__")
if node.target.__name__ == "special_logit"
)
for node in issue_nodes:
node.target = torch.special.logit

# Recompile the graph
gm.recompile()

return gm

def unstable_to_stable(self, gm):
methods = (
name
Expand Down
1 change: 1 addition & 0 deletions graph_net/torch/fx_graph_serialize_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ def serialize_graph_module_to_str(gm: torch.fx.GraphModule) -> str:
(r"torch\._C\._fft\.fft_irfft\(", "torch.fft.irfft("),
(r"torch\._C\._fft\.fft_rfft\(", "torch.fft.rfft("),
(r"torch\._C\._fft\.fft_fftn\(", "torch.fft.fftn("),
(r"torch\._C\._special\.special_logit\(", "torch.special.logit("),
# Add new rules to this list as needed
]
for pattern, repl in replacements:
Expand Down