github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-07 03:00:26 -04:00
Code Issues Packages Projects Releases Wiki Activity

ModNode.__mod__ simplification case

Browse Source
This commit is contained in:
chenyu
2024-01-02 11:49:40 -08:00
parent a53efa5eca
commit 73a42205a8
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
tinygrad/shape/symbolic.py
View File

@@ -198,8 +198,8 @@ class DivNode(OpNode):
class ModNode(OpNode):
def __mod__(self, b: Union[Node, int]):
if isinstance(b, Node) or isinstance(self.b, Node): return Node.__mod__(self, b)
return self.a % b if gcd(self.b, b) == b else Node.__mod__(self, b)
if isinstance(b, int) and isinstance(self.b, int) and b % self.b: return self.a % b
return Node.__mod__(self, b)
def __floordiv__(self, b: Union[Node, int], factoring_allowed=True):
if (self.b % b == 0): return (self.a//b) % (self.b//b) # put the div inside mod
return Node.__floordiv__(self, b, factoring_allowed)