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

Jump to bottom

[Inductor] aten.dist decomposition #105586

Closed
wants to merge 3 commits into from
Closed

[Inductor] aten.dist decomposition #105586

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
16fae64
[Inductor] aten.dist decomposition
yanboliang Jul 19, 2023
2e43de7
Update unit tests
yanboliang Jul 19, 2023
d9b2907
Add decomposition expect
yanboliang Jul 20, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 2 additions & 0 deletions test/expect/HasDecompTest.test_aten_core_operators.expect
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -171,6 +171,8 @@ aten::diagonal_scatter.out
aten::digamma
aten::digamma.out
aten::digamma_
aten::dist
aten::dist.out
aten::div.Scalar
aten::div.Scalar_mode
aten::div.Scalar_mode_out
Expand Down
2 changes: 0 additions & 2 deletions test/expect/HasDecompTest.test_has_decomposition.expect
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -714,8 +714,6 @@ aten::dequantize.tensors_out
aten::detach_
aten::detach_copy
aten::detach_copy.out
aten::dist
aten::dist.out
aten::embedding_renorm
aten::embedding_renorm.out
aten::embedding_renorm_
Expand Down
10 changes: 10 additions & 0 deletions test/inductor/test_torchinductor.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -976,6 +976,16 @@ def fn(a):

self.common(fn, (torch.randint(4, (4,)),))

def test_dist(self):
def fn(a, b):
return (
torch.dist(a, b),
torch.dist(a, b, p=1.2),
torch.dist(a.to(torch.bfloat16), b.to(torch.bfloat16)),
)

self.common(fn, (torch.randn(4, 4), torch.randn(4, 4)))

def test_arange1(self):
def fn(x):
rng1 = torch.arange(8 * 8, dtype=torch.float32, device=x.device).view(8, 8)
Expand Down
5 changes: 5 additions & 0 deletions torch/_decomp/decompositions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -630,6 +630,11 @@ def soft_margin_loss_backward(
return grad_input


@register_decomposition(aten.dist)
def dist(input: Tensor, other: Tensor, p: float = 2):
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This would be better suited for a ref, but fine

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This was done as part of a brief hackathon, so we were looking for the smallest tasks possible :P

return aten.norm(input - other, p=p)


@register_decomposition(aten._euclidean_dist)
def _euclidean_dist(x1: Tensor, x2: Tensor) -> Tensor:
x1_norm = x1.pow(2).sum(-1, True)
Expand Down
1 change: 1 addition & 0 deletions torch/_inductor/decomposition.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@
aten.bitwise_and_,
aten.bitwise_or_,
aten.clamp_min_,
aten.dist,
aten.empty_like,
aten.flip,
aten.lcm,
Expand Down
1 change: 0 additions & 1 deletion torch/_inductor/lowering.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1703,7 +1703,6 @@ def apply_constraint(arg, fx_arg):
make_fallback(aten.digamma, warn=False)
make_fallback(aten._efficientzerotensor)
make_fallback(aten._embedding_bag_per_sample_weights_backward)
make_fallback(aten.dist)
make_fallback(aten._efficientzerotensor)
make_fallback(aten._embedding_bag_per_sample_weights_backward)
make_fallback(aten.fractional_max_pool2d)
Expand Down