rangeify: fix abstractions2.py (#12386)

* rangeify: fix abstractions2.py

* tests

* lint

* only abstractions2

* base

.github/workflows/test.yml

@@ -540,6 +540,9 @@ jobs:
     - name: Test CPU=1 CPU_LLVM=1 RANGEIFY=1
       run: |
         CPU=1 CPU_LLVM=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_edgecases.py
+    - name: Test Docs RANGEIFY=1
+      run: |
+        RANGEIFY=1 python docs/abstractions2.py
     - name: Test const folding
       run: CPU=1 RANGEIFY=1 python3 -m pytest -n auto --durations 20 test/test_const_folding.py -k "not test_cast_padded and not TestReduceOpsConstFolding"
     # RANGEIFY=2 isn't supported

docs/abstractions2.py

@@ -80,7 +80,9 @@ print("******** third, the UOp ***********")
 
 from tinygrad.engine.realize import run_schedule
 from tinygrad.engine.schedule import create_schedule_with_vars
+from tinygrad.helpers import RANGEIFY
 from tinygrad.schedule.kernelize import get_kernelize_map
+from tinygrad.schedule.rangeify import get_rangeify_map
 
 # allocate some values + load in values
 a = UOp.new_buffer(DEVICE, 1, dtypes.int32)
@@ -93,10 +95,10 @@ out = a + b
 s = UOp(Ops.SINK, dtypes.void, (out,))
 
 # group the computation into kernels
-becomes_map = get_kernelize_map(s)
+becomes_map = get_rangeify_map(s) if RANGEIFY else get_kernelize_map(s)
 
 # the compute maps to an assign
-assign = becomes_map[a+b]
+assign = becomes_map[a+b].base
 
 # the first source is the output buffer (data)
 assert assign.src[0].op is Ops.BUFFER