diff --git a/tinygrad/codegen/opt/postrange.py b/tinygrad/codegen/opt/postrange.py index 9179cd3cb8..5614adec87 100644 --- a/tinygrad/codegen/opt/postrange.py +++ b/tinygrad/codegen/opt/postrange.py @@ -220,7 +220,7 @@ class Scheduler: def _apply_tc_opt(self, use_tensor_cores:int, axis:int, tc_select:int, opt_level:int) -> None|list[UOp]: if not (reduceops := self.reduceops): raise KernelOptError("no reduce ops for TensorCore") reduceop = reduceops[0] - if use_tensor_cores and reduceop is not None and reduceop.arg is Ops.ADD: + if use_tensor_cores and reduceop.arg is Ops.ADD: mul = reduceop.src[0] if reduceop.src[0].op is not Ops.CAST else reduceop.src[0].src[0] if mul.op is not Ops.MUL: return None in0, in1 = mul.src diff --git a/tinygrad/engine/jit.py b/tinygrad/engine/jit.py index 584826bbb7..a99cbe0a72 100644 --- a/tinygrad/engine/jit.py +++ b/tinygrad/engine/jit.py @@ -323,7 +323,6 @@ class TinyJit(Generic[ReturnType]): try: ret = self.fxn(*args, **kwargs) if len(params:=get_parameters(ret)): Tensor.realize(params[0], *params[1:]) - except Exception as e: raise e finally: capturing.clear() jit_cache = self._jit_cache del self._buffer_replace, self._jit_cache diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index e350a9a939..1ba458b893 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -308,12 +308,11 @@ class UOp(OpMixin, metaclass=UOpMetaClass): for er in self.ended_ranges: if er.op is Ops.RANGE: # if it's a single RANGE, we don't flow through it. - if er in ret: del ret[er] + ret.pop(er, None) else: # if it's not a RANGE, we include all ranges in srcs. # technically we shouldn't flow through these ranges either, but this is pre pm_add_control_flow so it's the same. - for s in er.ranges: - if s in ret: del ret[s] + for s in er.ranges: ret.pop(s, None) return ret @property