addressed the comments

Author: sharavak

lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp

@@ -10561,6 +10561,7 @@ namespace {
 //           | input,                  if target == 1
 // loss(x) = |
 //           | max(0, margin - input), if target == -1
+// target tensor may have values other than 1 and -1
 class DecomposeHingeEmbeddingLoss
     : public OpRewritePattern<AtenHingeEmbeddingLossOp> {
   using OpRewritePattern<AtenHingeEmbeddingLossOp>::OpRewritePattern;
@@ -10577,6 +10578,12 @@ class DecomposeHingeEmbeddingLoss
     auto targetTy = dyn_cast<ValueTensorType>(target.getType());
     if (!targetTy.hasDtype() || !targetTy.hasSizes())
       return rewriter.notifyMatchFailure(op, "target must have dtype and size");
+
+    int64_t reduction;
+    if (!matchPattern(op.getReduction(), m_TorchConstantInt(&reduction))) {
+      return rewriter.notifyMatchFailure(op,
+                                         "reduction should be a constant int!");
+    }
     auto resultTy = dyn_cast<ValueTensorType>(op.getType());
     Value minusOne = getConstantWithGivenDtypeAndValue(rewriter, loc, -1,
                                                        targetTy.getDtype());
@@ -10613,22 +10620,14 @@ class DecomposeHingeEmbeddingLoss
     // Add : outputMargin + outputSelf
     auto output = rewriter.create<AtenAddTensorOp>(loc, inputTy, outputMargin,
                                                    outputSelf, /*alpha=*/alpha);
-    int64_t reduction;
-    if (!matchPattern(op.getReduction(), m_TorchConstantInt(&reduction))) {
-      return rewriter.notifyMatchFailure(op,
-                                         "reduction should be a constant int!");
-    }
-    Value loss;
+    Value loss = output;
     Value none = rewriter.create<ConstantNoneOp>(loc);
     // reduction: mean
     if (reduction == 1) {
       loss = rewriter.create<AtenMeanOp>(loc, resultTy, output, none);
     } else if (reduction == 2) {
       // reduction: sum
       loss = rewriter.create<AtenSumOp>(loc, resultTy, output, none);
-    } else {
-      // reduction: none
-      loss = output;
     }
     rewriter.replaceOp(op, loss);
     return success();

projects/pt1/e2e_testing/xfail_sets.py

@@ -1797,7 +1797,7 @@
     "L1LossSumReductionModule_basic",
     "HingeEmbeddingLossReductionMeanModule_basic",
     "HingeEmbeddingLossReductionSumModule_basic",
-    "HingeEmbeddingLossWithoutReductionModule_basic",
+    "HingeEmbeddingLossReductionNoneModule_basic",
     "PixelShuffleModuleStaticRank3Int64_basic",
     "PixelShuffleModuleStaticRank4Float32_basic",
     "RandIntLowModule_basic",
@@ -2974,7 +2974,7 @@
     "HingeEmbeddingLossBasicModule_basic",
     "HingeEmbeddingLossReductionMeanModule_basic",
     "HingeEmbeddingLossReductionSumModule_basic",
-    "HingeEmbeddingLossWithoutReductionModule_basic",
+    "HingeEmbeddingLossReductionNoneModule_basic",
     "HstackBasicComplexModule_basic",
     "HstackBasicFloatModule_basic",
     "HstackBasicIntFloatModule_basic",
@@ -3988,7 +3988,7 @@
     "HingeEmbeddingLossBasicModule_basic",
     "HingeEmbeddingLossReductionMeanModule_basic",
     "HingeEmbeddingLossReductionSumModule_basic",
-    "HingeEmbeddingLossWithoutReductionModule_basic",
+    "HingeEmbeddingLossReductionNoneModule_basic",
     "Exp2StaticModule_basic",
     "ElementwiseRreluWithNoiseEvalModule_basic",
     "ElementwiseRreluWithNoiseEvalStaticModule_basic",

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

@@ -2189,8 +2189,7 @@ def aten〇nll_loss_backward〡shape(grad_output: List[int], self: List[int], ta
 def aten〇hinge_embedding_loss〡shape(self: List[int], target: List[int], margin: float = 1., reduction: int = 1) -> List[int]:
     if reduction in [1,2]:
         return []
-    else:
-        return upstream_shape_functions.unary(self)
+    return upstream_shape_functions.unary(self)
 
 # TODO: upstream this
 def aten〇mse_loss〡shape(self: List[int], target: List[int], reduction: int = 1) -> List[int]:

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

@@ -2486,8 +2486,8 @@ def __init__(self):
     @annotate_args(
         [
             None,
-            ([8, 1], torch.float32, True),
-            ([1, 1], torch.float32, True),
+            ([-1, -1], torch.float32, True),
+            ([-1, -1], torch.float32, True),
         ]
     )
     def forward(self, input, target):
@@ -2507,8 +2507,8 @@ def __init__(self):
     @annotate_args(
         [
             None,
-            ([2, 5], torch.float32, True),
-            ([1, 1], torch.float32, True),
+            ([-1, -1], torch.float32, True),
+            ([-1, -1], torch.float32, True),
         ]
     )
     def forward(self, input, target):
@@ -2520,24 +2520,24 @@ def HingeEmbeddingLossReductionSumModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(2, 5), tu.rand(1, 1))
 
 
-class HingeEmbeddingLossWithoutReductionModule(torch.nn.Module):
+class HingeEmbeddingLossReductionNoneModule(torch.nn.Module):
     def __init__(self):
         super().__init__()
 
     @export
     @annotate_args(
         [
             None,
-            ([8, 5], torch.float32, True),
-            ([1], torch.float32, True),
+            ([-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: HingeEmbeddingLossWithoutReductionModule())
-def HingeEmbeddingLossWithoutReductionModule_basic(module, tu: TestUtils):
+@register_test_case(module_factory=lambda: HingeEmbeddingLossReductionNoneModule())
+def HingeEmbeddingLossReductionNoneModule_basic(module, tu: TestUtils):
     module.forward(tu.rand(8, 5), tu.rand(1))