From fa88e1d0d02beda195c5679adf3ecfd84b0c53a2 Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Wed, 28 Feb 2024 08:15:01 -0500
Subject: [PATCH] cleanup lazy reduce (#3517)

* cleanup lazy reduce

removed useless assert now arg is axis and cleaned split logic

* stride can be symbolic with int shape
---
 tinygrad/lazy.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py
index 7632de4a09..b1ad3effb1 100644
--- a/tinygrad/lazy.py
+++ b/tinygrad/lazy.py
@@ -131,11 +131,11 @@ class LazyBuffer:
     new_shape = tuple(1 if i in axis else s for i,s in enumerate(self.shape))
     # TODO: this logic should move to the scheduler
     if self.size == 0 and 0 not in new_shape: return self.const({ReduceOps.SUM: 0.0, ReduceOps.MAX: -math.inf}[op], new_shape)
-    assert len(self.shape)==len(new_shape) and all(ns in (1,s) for s,ns in zip(self.shape,new_shape)), f"not a contraction {self.shape=} {new_shape=}"
     # TODO: can we split symbolic shape if the reduce axis is not symbolic?
     if not all_int(self.shape) or (0 in self.shape) or prod(self.shape) // prod(new_shape) < getenv("REDUCEOP_SPLIT_THRESHOLD", 32768):
       return self._reduce_op(op, axis)
-    heuristic, divisor, dim_to_split = max(((divisor := math.gcd(256, old))/(stride or math.inf), divisor, i) for i, (old, new, stride) in enumerate(zip(self.shape, new_shape, self.st.real_strides())) if old != new) # type: ignore  # noqa: E501
+    heuristic, divisor, dim_to_split = max(((divisor := math.gcd(256, s))/(st or math.inf), divisor, i) for i,(s,st) in \
+                                           enumerate(zip(self.shape, self.st.real_strides())) if i in axis and (st is None or isinstance(st, int)))
     if divisor < 16 or heuristic < 0.1: return self._reduce_op(op, axis)
     # choose largest divisor (>=16) to split on, penalize large strides
     def splitted_shape(dim_aft_div):