diff --git a/test/test_schedule.py b/test/test_schedule.py index e8fe0ab5f7..7d48eee154 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -143,13 +143,14 @@ class TestSchedule(unittest.TestCase): root = root + functools.reduce(lambda a,b:a+b, bufs[i:i+X]) self.assertEqual(root.item(), sum(range(N))) - @unittest.expectedFailure # TODO: failing because of can_chase - def test_indexing_scalars_multiple_dims(self): - X = Tensor.randn(2, 3).realize() - xt = X[Tensor(0)][Tensor(1)] + @given(strat.sampled_from(range(2,4)), strat.sampled_from(range(2,4)), strat.sampled_from(range(0,4)), strat.sampled_from(range(0,4))) + def test_indexing_scalars(self, x, y, a, b): + assume(a dict[UOp, None]: recursive_group(r, unwrap(r.st), r, children, realizes, reduce_for_op, group, cache={}) # max one reduceop per kernel can_chase = all(tr not in reduce_for_op for tr in group) + for u in r.toposort(gate=lambda u: u not in realizes): + if u.op is Ops.REDUCE_AXIS and u.src[0].base.op is Ops.CONST: + can_chase = False + break # TODO: forced_realize exists because the scheduler is incapable of checking for self-contained DAGs forced_realize = r in group # can only have one output