Fix xlogy returning 0 instead of NaN when self is 0 and other is NaN

onnxscript/function_libs/torch_lib/ops/special.py

@@ -375,8 +375,8 @@ def aten_special_xlogy(self: TFloat, other: TFloat) -> TFloat:
     nans = op.IsNaN(other)
     zeros = op.Equal(self, 0)
     xlogy = op.Mul(self, op.Log(other))
-    xlogy_with_nans = op.Where(nans, other, xlogy)
-    return op.Where(zeros, self, xlogy_with_nans)
+    xlogy_with_zeros = op.Where(zeros, self, xlogy)
+    return op.Where(nans, other, xlogy_with_zeros)
 
 
 def aten_special_zeta(self: TensorType, other: TensorType) -> TensorType:

tests/function_libs/torch_lib/e2e_ops_tests.py

@@ -157,6 +157,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         )
         _testing.assert_onnx_program(onnx_program)
 
+    def test_xlogy_nan_other_takes_precedence_over_zero_self(self):
+        # torch.xlogy returns NaN wherever other is NaN, even where self is 0, so
+        # the NaN case has to win over the zero case.
+        class Model(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.xlogy(x, y)
+
+        nan = math.nan
+        x = torch.tensor([0.0, 0.0, 1.0, 2.0, 0.0])
+        y = torch.tensor([nan, 2.0, nan, 3.0, 0.0])
+
+        onnx_program = torch.onnx.export(Model(), (x, y), dynamo=True, optimize=False)
+        _testing.assert_onnx_program(onnx_program)
+
     def test_rand_like_memory_format(self):
         # These random *_like ops are non-deterministic, so assert the export
         # succeeds rather than comparing values (see issue #3002).