[TORCH] Add support for aten.hinge_embedding_loss Op #4227

sharavak · 2025-06-10T14:21:53Z

Decomposed hinge_embedding_loss op into Aten ops.
Added test cases in the e2e part.

This implementation addresses and closes #4222

sharavak · 2025-06-10T15:40:37Z

@stellaraccident , @vivekkhandelwal1, @zjgarvey I’d be grateful if any of you could take a look at this PR. Your feedback would be greatly appreciated!

projects/pt1/python/torch_mlir_e2e_test/test_suite/reduction.py

projects/pt1/python/torch_mlir/jit_ir_importer/build_tools/abstract_interp_lib_gen.py

vivekkhandelwal1 · 2025-06-16T12:22:30Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    // Compute mask: target != 1
+    auto targetNotOne =
+        rewriter.create<AtenNeScalarOp>(loc, boolType, target, one);
+    // If target != 1 use marginClamp otherwise 0.
+    auto outputMargin = rewriter.create<AtenWhereScalarOtherOp>(
+        loc, inputTy, targetNotOne, marginClamp, zero);
+    // Compute mask: target != -1
+    auto targetNotMinusOne =
+        rewriter.create<AtenNeScalarOp>(loc, boolType, target, minusOne);
+    // If target != -1 use the original input. Otherwise 0.
+    auto outputSelf = rewriter.create<AtenWhereScalarOtherOp>(
+        loc, inputTy, targetNotMinusOne, input, zero);
+    // Add : outputMargin + outputSelf
+    auto output = rewriter.create<AtenAddTensorOp>(loc, inputTy, outputMargin,
+                                                   outputSelf, /*alpha=*/alpha);


Instead of doing all this, you can just do:

auto result = rewriter.create<AtenWhereScalarOtherOp>( loc, inputTy, targetNotOne, marginClamp, input);

@vivekkhandelwal1 Thanks for the suggestion, I did try the simplified version initially, but it caused numerical validation errors in some test cases. This happens because the target tensor can sometimes have values other than just -1 and 1.

To handle this properly and stay consistent with PyTorch's semantics, I decided to explicitly check for both target == 1 and target == -1. This way, the behavior stays correct even if target have values other than just -1 and 1.

Eg:

import torch input=torch.randn(2,3) target=torch.randn(2,3) torch.hinge_embedding_loss(input,target) Output: tensor([[1.1361, 1.0000, 1.0000], [1.4880, 1.1624, 1.0000]])

This happens because the target tensor can sometimes have values other than just -1 and 1.

In what cases and how? Since the definition says that it can contain only -1 and 1.

Thanks for the reply @vivekkhandelwal1, I got reference from the Pytorch native implementation https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/Loss.cpp#L182. While the official definition of hinge_embedding_loss states that the target should contain only -1 and 1, the native implementation doesn’t enforce this restriction and handles arbitrary values using at::where(target != 1, ) and at::where(target != -1,). So to stay with the Pytorch behaviour, I followed the same logic.

Okay. In that case, it would be better to add this link and the justification in the comment. So that, the future users/contributors may not get confused.

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

sharavak · 2025-06-17T17:52:07Z

@vivekkhandelwal1 Thanks a lot for the feedback. I've updated the code.

vivekkhandelwal1 · 2025-06-30T05:44:27Z

projects/pt1/python/torch_mlir_e2e_test/test_suite/reduction.py

+def HingeEmbeddingLossReductionMeanModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(8, 1), tu.rand(1, 1))
+
+
+class HingeEmbeddingLossReductionSumModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1], torch.float32, True),
+            ([-1, -1], torch.float32, True),
+        ]
+    )
+    def forward(self, input, target):
+        return torch.ops.aten.hinge_embedding_loss(input, target, reduction=2)
+
+
+@register_test_case(module_factory=lambda: HingeEmbeddingLossReductionSumModule())
+def HingeEmbeddingLossReductionSumModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(2, 5), tu.rand(1, 1))
+
+
+class HingeEmbeddingLossReductionNoneModule(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    @export
+    @annotate_args(
+        [
+            None,
+            ([-1, -1], torch.float32, True),
+            ([-1], torch.float32, True),
+        ]
+    )
+    def forward(self, input, target):
+        return torch.ops.aten.hinge_embedding_loss(input, target, margin=1.0)
+
+
+@register_test_case(module_factory=lambda: HingeEmbeddingLossReductionNoneModule())
+def HingeEmbeddingLossReductionNoneModule_basic(module, tu: TestUtils):
+    module.forward(tu.rand(8, 5), tu.rand(1))


All 3 of these tests have only a single value as a target; as a result, not all the paths of the lowering will get tested. Ideally, for testing purposes, the target tensor should contain a mix of -1 and 1 values. Also, how do you make sure that the target tensor contains only valid values?

Thanks for the review @vivekkhandelwal1
I will update the target tensor.

Also, how do you make sure that the target tensor contains only valid values?

For this question, I’ve addressed that in my earlier comment here: #4227 (comment). @vivekkhandelwal1.

vivekkhandelwal1 · 2025-06-30T06:08:06Z

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

+    // Compute mask: target != 1
+    auto targetNotOne =
+        rewriter.create<AtenNeScalarOp>(loc, boolType, target, one);
+    // If target != 1 use marginClamp otherwise 0.
+    auto outputMargin = rewriter.create<AtenWhereScalarOtherOp>(
+        loc, inputTy, targetNotOne, marginClamp, zero);
+    // Compute mask: target != -1
+    auto targetNotMinusOne =
+        rewriter.create<AtenNeScalarOp>(loc, boolType, target, minusOne);
+    // If target != -1 use the original input. Otherwise 0.
+    auto outputSelf = rewriter.create<AtenWhereScalarOtherOp>(
+        loc, inputTy, targetNotMinusOne, input, zero);
+    // Add : outputMargin + outputSelf
+    auto output = rewriter.create<AtenAddTensorOp>(loc, inputTy, outputMargin,
+                                                   outputSelf, /*alpha=*/alpha);


This happens because the target tensor can sometimes have values other than just -1 and 1.

In what cases and how? Since the definition says that it can contain only -1 and 1.

sharavak force-pushed the hinge_embedding_loss branch 2 times, most recently from 9d5b896 to ddf112a Compare June 10, 2025 14:40

sharavak marked this pull request as ready for review June 10, 2025 14:41

vivekkhandelwal1 self-requested a review June 16, 2025 12:13

vivekkhandelwal1 requested changes Jun 16, 2025

View reviewed changes

sharavak added 2 commits June 17, 2025 17:11

add the code for aten.hinge_embedding_loss op

712b20f

addressed the comments

8d8c30b

sharavak force-pushed the hinge_embedding_loss branch from ddf112a to 8d8c30b Compare June 17, 2025 17:13

sharavak requested a review from vivekkhandelwal1 June 17, 2025 17:48

sharavak added 2 commits June 19, 2025 10:27

Merge branch 'main' into hinge_embedding_loss

a6b1da8

Merge branch 'main' into hinge_embedding_loss

9aa8605

vivekkhandelwal1 requested changes Jun 30, 2025

View reviewed changes

updated the test case

232c2b8

sharavak requested a review from vivekkhandelwal1 July 2, 2025 17:14

sharavak and others added 2 commits July 9, 2025 21:15

Merge branch 'main' into hinge_embedding_loss

5d960ff

Fix pre-commit hook issues (trailing whitespace)

c70507b

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[TORCH] Add support for aten.hinge_embedding_loss Op #4227

[TORCH] Add support for aten.hinge_embedding_loss Op #4227

Uh oh!

sharavak commented Jun 10, 2025

Uh oh!

sharavak commented Jun 10, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vivekkhandelwal1 Jun 16, 2025

Uh oh!

sharavak Jun 17, 2025 •

edited

Loading

Uh oh!

vivekkhandelwal1 Jun 30, 2025

Uh oh!

sharavak Jul 1, 2025 •

edited

Loading

Uh oh!

vivekkhandelwal1 Jul 15, 2025

Uh oh!

Uh oh!

Uh oh!

sharavak commented Jun 17, 2025 •

edited

Loading

Uh oh!

vivekkhandelwal1 Jun 30, 2025

Uh oh!

sharavak Jul 1, 2025 •

edited

Loading

Uh oh!

vivekkhandelwal1 Jun 30, 2025

Uh oh!

Uh oh!

[TORCH] Add support for aten.hinge_embedding_loss Op #4227

Are you sure you want to change the base?

[TORCH] Add support for aten.hinge_embedding_loss Op #4227

Uh oh!

Conversation

sharavak commented Jun 10, 2025

Uh oh!

sharavak commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

vivekkhandelwal1 Jun 16, 2025

Choose a reason for hiding this comment

Uh oh!

sharavak Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vivekkhandelwal1 Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

sharavak Jul 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vivekkhandelwal1 Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sharavak commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

vivekkhandelwal1 Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

sharavak Jul 1, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

vivekkhandelwal1 Jun 30, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

sharavak commented Jun 10, 2025 •

edited

Loading

sharavak Jun 17, 2025 •

edited

Loading

sharavak Jul 1, 2025 •

edited

Loading

sharavak commented Jun 17, 2025 •

edited

Loading

sharavak Jul 1, 2025 •

edited

Loading