tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-09-02 20:46:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

support float mod (#9306)

Browse Source 
also added spec check on Ops.MOD to be ints only
This commit is contained in:
chenyu
2025-02-28 13:33:58 -05:00
committed by GitHub
parent 90808e2dd0
commit 7c7db78feb
3 changed files with 15 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_ops.py
+13 -5
View File
@@ -582,11 +582,19 @@ class TestOps(unittest.TestCase):
helper_test_op([()], lambda x: 2/x)
def test_mod(self):
helper_test_op(None, lambda x,y: x%y, Tensor.mod, forward_only=True, vals=[[-4, 7, 5, 4, -7, 8], [2, -3, 8, -2, 3, 5]])
helper_test_op(None, lambda x,y: x%y, forward_only=True, vals=[[-4, 7, 5, 4, -7, 8], [2, -3, 8, -2, 3, 5]])
helper_test_op(None, lambda x: x%2, forward_only=True, vals=[[-4, 7, 5, 4, -7, 8]])
helper_test_op(None, lambda x: x%3, forward_only=True, vals=[[-4, 7, 5, 4, -7, 8]])
helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[[-4, 7, 5, 4, -7, 8]])
a = [-4, 7, 5, 4, -7, 8]
b = [2, -3, 8, -2, 3, 5]
for float_a in [True, False]:
for float_b in [True, False]:
va = [float(ai) for ai in a] if float_a else a
vb = [float(bi) for bi in b] if float_b else b
helper_test_op(None, lambda x,y: x%y, Tensor.mod, forward_only=True, vals=[va, vb])
helper_test_op(None, lambda x,y: x%y, forward_only=True, vals=[va, vb])
helper_test_op(None, lambda x: x%2, forward_only=True, vals=[va])
helper_test_op(None, lambda x: x%3, forward_only=True, vals=[va])
helper_test_op(None, lambda x: x%3.5, forward_only=True, vals=[va])
helper_test_op(None, lambda x: 100%x, forward_only=True, vals=[va])
helper_test_op(None, lambda x: 100.5%x, forward_only=True, vals=[va])
def test_mul_naninf(self):
helper_test_op([(45,65)], lambda x: x*math.inf)
tinygrad/spec.py
+1 -1
View File
@@ -91,7 +91,7 @@ spec = PatternMatcher([
(UPat((Ops.CMPLT, Ops.CMPNE), dtype=dtypes.bool, src=(UPat.var("x"), UPat.var("y"))), lambda x,y: x.dtype.base == y.dtype.base),
# and SHL/SHR, the shift distance can be an int
(UPat((Ops.SHL, Ops.SHR), src=(UPat.var("x"), UPat.var("y")), name="a"), lambda a,x,y: a.dtype == x.dtype and y.dtype in (x.dtype, dtypes.uint)),
(UPat(Ops.IDIV, name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
(UPat((Ops.IDIV, Ops.MOD), name="x"), lambda x: None if dtypes.is_int(x.dtype) else False),
(UPat(GroupOp.ALU, name="x"), lambda x: all(x.dtype.base == y.dtype.base for y in x.src)),
(UPat(Ops.ASSIGN, src=(UPat((Ops.DEFINE_ACC, Ops.DEFINE_GLOBAL)), UPat())), lambda: True),
tinygrad/tensor.py
+1 -1
View File
@@ -3289,7 +3289,7 @@ class Tensor(SimpleMathTrait):
```
"""
a, b = self._broadcasted(x, reverse)
return (r := a._apply_uop(UOp.mod, b)) + b * (((r < 0) & (b > 0)) | ((r > 0) & (b < 0)))
return a - a.div(b, rounding_mode="floor") * b
def bitwise_xor(self, x:Union[Tensor, ConstType], reverse=False) -> Tensor:
"""