From 954a2fef75fc2867d308b2d52a339227ff1b27c8 Mon Sep 17 00:00:00 2001
From: George Hotz <geohot@gmail.com>
Date: Mon, 18 Dec 2023 16:33:22 -0800
Subject: [PATCH] hotfix: add JITGRAPH and invert sints

---
 tinygrad/jit.py                | 5 +++--
 tinygrad/shape/shapetracker.py | 2 +-
 tinygrad/shape/view.py         | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/tinygrad/jit.py b/tinygrad/jit.py
index 7caf04a0f0..85ea4f2721 100644
--- a/tinygrad/jit.py
+++ b/tinygrad/jit.py
@@ -1,7 +1,7 @@
 from __future__ import annotations
 from typing import Callable, List, Tuple, Dict, cast, Union, Optional, TypeVar, Generic
 import functools, itertools, operator
-from tinygrad.helpers import DEBUG, DType, merge_dicts, getenv, all_int
+from tinygrad.helpers import DEBUG, DType, merge_dicts, getenv, all_int, Context, GRAPH
 from tinygrad.device import Device, JITRunner, CompiledASTRunner, Buffer
 from tinygrad.tensor import Tensor
 from tinygrad.shape.shapetracker import ShapeTracker
@@ -71,7 +71,8 @@ class TinyJit(Generic[ReturnType]):
       # jit capture
       self.expected_vals, self.expected_name_sts_dtype = expected_vals, expected_name_sts_dtype
       CacheCollector.start(var_vals)
-      self.ret = self.fxn(*args, **kwargs)
+      with Context(GRAPH=getenv("JITGRAPH", GRAPH.value)):
+        self.ret = self.fxn(*args, **kwargs)
       self.jit_cache = CacheCollector.finish()
       assert len(self.jit_cache) != 0, "didn't JIT anything!"
       assert len(set(get_input_replace(self.jit_cache, input_rawbuffers).values())) == len(input_rawbuffers), "some input tensors not found"
diff --git a/tinygrad/shape/shapetracker.py b/tinygrad/shape/shapetracker.py
index 330f1ed433..6598fbfb13 100644
--- a/tinygrad/shape/shapetracker.py
+++ b/tinygrad/shape/shapetracker.py
@@ -58,7 +58,7 @@ class ShapeTracker:
 
   def __add__(self, st:ShapeTracker) -> ShapeTracker: return ShapeTracker(self.views + st.views).simplify()
 
-  def invert(self, out_shape:Tuple[int, ...]) -> Optional[ShapeTracker]:
+  def invert(self, out_shape:Tuple[sint, ...]) -> Optional[ShapeTracker]:
     ret = tuple(v.invert(s) for v,s in zip(self.views[::-1], [x.shape for x in self.views[::-1][1:]]+[out_shape]))
     return ShapeTracker(cast(Tuple[View, ...], ret)) if all(x is not None for x in ret) else None
 
diff --git a/tinygrad/shape/view.py b/tinygrad/shape/view.py
index ff378e56ba..8ebd420c96 100644
--- a/tinygrad/shape/view.py
+++ b/tinygrad/shape/view.py
@@ -99,7 +99,7 @@ class View:
     return View.create(new_shape, new_strides, new_offset, new_mask)
 
   @functools.lru_cache(maxsize=None)  # pylint: disable=method-cache-max-size-none
-  def invert(self, out_shape:Tuple[int, ...]) -> Optional[View]:
+  def invert(self, out_shape:Tuple[sint, ...]) -> Optional[View]:
     ret = self.shrink(self.mask) if self.mask else self
     if prod(ret.shape) != prod(out_shape): return None   # don't support shrink, expand, or stride != (-1, 1)
     ret = cast(View, ret.reshape(tuple(s for s in ret.shape if s != 1)))  # removing ones will never be an issue