make maximum split the grad like torch when equal (#738)

* make maximum split grad

* added test for maximum split grad when equal

* minor expr simplification

* (2-eq)/2 only once

* update test bc one more sum output child stays

test/external/external_test_opt.py

@@ -190,7 +190,7 @@ class TestOpt(unittest.TestCase):
       # TODO: this should be 4, but the sum output child stays around
       # with pushing_permutes it can be 3
       # TODO: broken with optim fixes
-      assert len(GlobalCounters.cache) in [4,5,6,7], f"optimizer didn't fold conv-backward SGD, got {len(GlobalCounters.cache)}"
+      assert len(GlobalCounters.cache) in [4,5,6,7,8], f"optimizer didn't fold conv-backward SGD, got {len(GlobalCounters.cache)}"
     Tensor.training = False
 
   def test_fold_conv_batchnorm_sgd(self):

test/test_ops.py

@@ -91,6 +91,7 @@ class TestOps(unittest.TestCase):
 
   def test_maximum(self):
     helper_test_op([(45,65), (45,65)], torch.maximum, Tensor.maximum)
+    helper_test_op(None, torch.maximum, Tensor.maximum, vals=[[1., 2., 3., 4.], [1., 2., 3., 4.]])
   def test_minimum(self):
     helper_test_op([(45,65), (45,65)], torch.minimum, Tensor.minimum)
   def test_add(self):

tinygrad/mlops.py

@@ -67,14 +67,17 @@ class Equal(Function):
 
 class Maximum(Function):
   def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer:
-    self.y, self.ret = y, x.binary_op(BinaryOps.MAX, y)
+    self.x, self.y = x, y
+    self.ret = x.binary_op(BinaryOps.MAX, y)
     return self.ret
 
   def backward(self, grad_output):
     mask = self.y.binary_op(BinaryOps.CMPEQ, self.ret)
-    # TODO: if they are equal, do they split the gradient?
-    return grad_output.binary_op(BinaryOps.MUL, mask.const_like(1).binary_op(BinaryOps.SUB, mask)) if self.needs_input_grad[0] else None, \
-           grad_output.binary_op(BinaryOps.MUL, mask) if self.needs_input_grad[1] else None
+    eq = self.x.binary_op(BinaryOps.CMPEQ, self.y)
+    splitter = eq.const_like(2).binary_op(BinaryOps.SUB, eq).binary_op(BinaryOps.DIV, eq.const_like(2))
+
+    return grad_output.binary_op(BinaryOps.MUL, mask.const_like(1).binary_op(BinaryOps.SUB, mask).binary_op(BinaryOps.ADD, eq)).binary_op(BinaryOps.MUL, splitter) if self.needs_input_grad[0] else None, \
+           grad_output.binary_op(BinaryOps.MUL, mask).binary_op(BinaryOps.MUL, splitter) if self.needs_input_grad[1] else None
 
 class Add(Function):
   def forward(self, x:LazyBuffer, y:LazyBuffer) -> LazyBuffer: