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

[Relax][PyTorch] Cleanup Tensor Manipulation and Creation op converters #17376

Merged
merged 29 commits into from
Sep 16, 2024
Merged

[Relax][PyTorch] Cleanup Tensor Manipulation and Creation op converters #17376

Changes from 1 commit
Commits
Show all changes
29 commits
Select commit Hold shift + click to select a range
5be96be
cleanup `_cat()`
mshr-h Sep 14, 2024
363a7e6
cleanup `_cumsum()`
mshr-h Sep 14, 2024
16631fe
cleanup `_expand()`
mshr-h Sep 14, 2024
9c6a9b6
cleanup `_flatten()`
mshr-h Sep 14, 2024
b96f594
cleanup `_permute()`
mshr-h Sep 14, 2024
6cbd8aa
cleanup `_repeat()`
mshr-h Sep 14, 2024
3025d2c
cleanup `_reshape()`
mshr-h Sep 14, 2024
fe05adf
cleanup `_size()`
mshr-h Sep 14, 2024
a387888
cleanup `_split()`
mshr-h Sep 14, 2024
4c51b88
cleanup `_squeeze()`
mshr-h Sep 14, 2024
4416225
cleanup `_tile()`
mshr-h Sep 14, 2024
1cfc75d
cleanup `_transpose()`
mshr-h Sep 14, 2024
c9845aa
cleanup `chunk()`
mshr-h Sep 14, 2024
13f6657
cleanup `_arange()`
mshr-h Sep 14, 2024
2716433
cleanup `_empty()`
mshr-h Sep 14, 2024
b516828
cleanup `_inplace_fill()`
mshr-h Sep 14, 2024
f657b9c
cleanup `_full()`
mshr-h Sep 14, 2024
a21b1f5
cleanup `_index_select()`
mshr-h Sep 14, 2024
64a7edd
cleanup `_inplace_masked_fill()`
mshr-h Sep 14, 2024
3d90ce1
cleanup `_masked_fill()`
mshr-h Sep 14, 2024
98ae305
cleanup `_new_ones()`
mshr-h Sep 14, 2024
63512f8
cleanup `_ones()`
mshr-h Sep 14, 2024
0305588
cleanup `_tensor()`
mshr-h Sep 14, 2024
5bd647f
`_inplace_tril_triu()` is an unary op
mshr-h Sep 14, 2024
e35093e
`_batch_norm_2d()` is a nn ops
mshr-h Sep 14, 2024
49e5576
`_interpolate()` is a nn ops
mshr-h Sep 14, 2024
669ca0f
`_cross_entropy()` is a nn ops
mshr-h Sep 14, 2024
7521211
chore
mshr-h Sep 14, 2024
f78e9cd
fix tensor size
mshr-h Sep 16, 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
Prev Previous commit
Next Next commit
cleanup _ones()
  • Loading branch information
@mshr-h
mshr-h committed Sep 15, 2024
commit 63512f8c09016a28b5a504783b10bf0f9bf490d3
37 changes: 16 additions & 21 deletions python/tvm/relax/frontend/torch/fx_translator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1154,6 +1154,22 @@ def _new_ones(self, node: fx.Node) -> relax.Var:
)
)

def _ones(self, node: fx.Node) -> relax.Var:
import torch

args = self.retrieve_args(node)
size = relax.ShapeExpr((args[0],) if isinstance(args[0], (list, tuple)) else args[0])
dtype = self._convert_data_type(
node.kwargs.get("dtype", torch.get_default_dtype()), self.env
)
return self.block_builder.emit(
relax.op.full(
size,
relax.const(1, dtype),
dtype,
)
)

########## DataType ##########

def _float(self, node: fx.Node) -> relax.Var:
Expand Down Expand Up @@ -1204,27 +1220,6 @@ def convert(node: fx.Node) -> relax.Var:

return convert

def _ones(self, node: fx.Node) -> relax.Var:
import torch

args = self.retrieve_args(node)
size = args[0]
if not isinstance(size, (list, tuple)):
size = (size,)
size = relax.ShapeExpr(size)
dtype = (
TorchFXImporter._convert_data_type(str(node.kwargs["dtype"]), self.env)
if "dtype" in node.kwargs
else TorchFXImporter._convert_data_type(torch.get_default_dtype(), self.env)
)
return self.block_builder.emit(
relax.op.full(
size,
relax.const(1, dtype),
dtype,
)
)

########## Manipulation ##########

########## Neural Network ##########
Expand Down