fix: split fmod into separate tensor/scalar entries, fix typo in test name

Author: factnn

src/flag_gems/__init__.py

@@ -215,10 +215,10 @@ def torch_ge(v):
     ("floor_divide_.Tensor", floor_divide_),
     ("fmin", fmin),
     ("fmin.out", fmin_out),
-    ("fmod.Scalar", fmod),
-    ("fmod.Tensor", fmod),
-    ("fmod_.Scalar", fmod_),
-    ("fmod_.Tensor", fmod_),
+    ("fmod.Scalar", fmod_scalar),
+    ("fmod.Tensor", fmod_tensor),
+    ("fmod_.Scalar", fmod_scalar_),
+    ("fmod_.Tensor", fmod_tensor_),
     ("full", full),
     ("full_like", full_like),
     ("gather", gather),

src/flag_gems/ops/__init__.py

@@ -125,7 +125,7 @@
 from flag_gems.ops.flip import flip
 from flag_gems.ops.floor_ import floor_
 from flag_gems.ops.fmin import fmin, fmin_out
-from flag_gems.ops.fmod import fmod, fmod_
+from flag_gems.ops.fmod import fmod_scalar, fmod_scalar_, fmod_tensor, fmod_tensor_
 from flag_gems.ops.full import full
 from flag_gems.ops.full_like import full_like
 from flag_gems.ops.gather import gather, gather_backward
@@ -480,8 +480,10 @@
     "floor_divide_",
     "fmin",
     "fmin_out",
-    "fmod",
-    "fmod_",
+    "fmod_scalar",
+    "fmod_scalar_",
+    "fmod_tensor",
+    "fmod_tensor_",
     "full",
     "full_like",
     "gather",

src/flag_gems/ops/fmod.py

@@ -1,6 +1,5 @@
 import logging
 
-import torch
 import triton
 import triton.language as tl
 
@@ -32,20 +31,21 @@ def fmod_func_tensor_scalar(x, y):
     return result.to(dtype)
 
 
-def fmod(A, B):
-    logger.debug("GEMS FMOD")
-    if isinstance(A, torch.Tensor) and isinstance(B, torch.Tensor):
-        return fmod_func(A, B)
-    elif isinstance(A, torch.Tensor):
-        return fmod_func_tensor_scalar(A, B)
-    else:
-        # Both scalar - fallback to PyTorch
-        return torch.fmod(torch.tensor(A), B)
+def fmod_tensor(A, B):
+    logger.debug("GEMS FMOD TENSOR")
+    return fmod_func(A, B)
 
 
-def fmod_(A, B):
-    logger.debug("GEMS FMOD_")
-    if isinstance(B, torch.Tensor):
-        return fmod_func(A, B, out0=A)
-    else:
-        return fmod_func_tensor_scalar(A, B, out0=A)
+def fmod_scalar(A, B):
+    logger.debug("GEMS FMOD SCALAR")
+    return fmod_func_tensor_scalar(A, B)
+
+
+def fmod_tensor_(A, B):
+    logger.debug("GEMS FMOD_ TENSOR")
+    return fmod_func(A, B, out0=A)
+
+
+def fmod_scalar_(A, B):
+    logger.debug("GEMS FMOD_ SCALAR")
+    return fmod_func_tensor_scalar(A, B, out0=A)

tests/test_fmod.py

@@ -61,7 +61,7 @@ def test_fmod_tensor_inplace(shape, dtype):
 @pytest.mark.parametrize("shape", POINTWISE_SHAPES)
 @pytest.mark.parametrize("scalar", SCALARS)
 @pytest.mark.parametrize("dtype", FLOAT_DTYPES)
-def testy_fmod_scalar_inplace(shape, scalar, dtype):
+def test_fmod_scalar_inplace(shape, scalar, dtype):
     inp1 = torch.randn(shape, dtype=dtype, device=flag_gems.device)
     inp2 = scalar if scalar != 0 else 1.0
     ref_inp1 = to_reference(inp1.clone(), True)