Enable threefry_doesnt_use_long test on rangeify (#12229)

* dont bufferize rangeify

* enable doesnt_use_long test
This commit is contained in:
Sieds Lykles
2025-09-18 01:58:34 +02:00
committed by GitHub
parent 3c5b8bf50c
commit 812f485cd7
2 changed files with 3 additions and 2 deletions
+1 -1
View File
@@ -525,7 +525,7 @@ jobs:
# test_instancenorm_3d is very slow
run: |
CPU=1 CPU_LLVM=0 RANGEIFY=1 python3 -m pytest -n auto --durations 20 \
-k "not test_threefry_doesnt_use_long and not test_embedding and not test_load_state_dict_sharded_model_dict_same_axis and not test_instancenorm_3d" \
-k "not test_embedding and not test_load_state_dict_sharded_model_dict_same_axis and not test_instancenorm_3d" \
test/test_tiny.py test/test_rangeify.py test/test_ops.py test/test_tensor_variable.py \
test/test_outerworld_range.py test/test_sample.py test/test_randomness.py test/test_nn.py test/test_arange.py
- name: Test const folding
+2 -1
View File
@@ -376,7 +376,8 @@ def remove_bufferize(src:UOp, buf:UOp, idx:UOp):
# here is where we compute the cost
# for now just no REDUCE, COPY, or ASSIGN
ran = src.toposort(gate=lambda x: x.op not in {Ops.INDEX})
if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.ASSIGN} for x in ran): return None
# we don't want to bufferize threefry, also causes problems because not all platforms support long
if any(x.op in {Ops.REDUCE, Ops.COPY, Ops.ASSIGN} for x in ran) and src.op is not Ops.THREEFRY: return None
# simple, matching old behavior
#if src.op is not Ops.INDEX: return None