Skip to content
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

Optimize _split_xla_args_tensor_sym_constant #7900

Merged
merged 3 commits into from
Aug 22, 2024
Merged

Optimize _split_xla_args_tensor_sym_constant #7900

Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
090d441
Optimize _split_xla_args_tensor_sym_constant
JackCaoG Aug 22, 2024
8990609
remove unneeded time
JackCaoG Aug 22, 2024
6517ab1
remove unneeded nonlocal
JackCaoG Aug 22, 2024
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
5 changes: 3 additions & 2 deletions torch_xla/core/dynamo_bridge.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@
from typing import Any, Dict, List, Set, Tuple, Union
from numbers import Number
from contextlib import contextmanager
from collections import deque

import torch
from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner
Expand Down Expand Up @@ -183,7 +184,7 @@ def _maybe_move_tensors_to_device(tensors: tuple,


def _split_xla_args_tensor_sym_constant(args):
tensors = []
tensors = deque(maxlen=len(args))
constants = []
for arg in args:
if isinstance(arg, torch.Tensor):
Expand All @@ -193,7 +194,7 @@ def _split_xla_args_tensor_sym_constant(args):
else:
assert False, "argument to the xla dynamo bridge should be either a number or tensor"

return tuple(tensors), tuple(constants)
return tensors, tuple(constants)


class Deduper:
Expand Down
Loading