def test_arange_simple():
t = randn((3, 4, 5))
f = Tensor(t)["i", "j"]
assert_close(f, f(i=f.new_arange("i", 3)))
assert_close(f, f(j=f.new_arange("j", 4)))
assert_close(f, f(i=f.new_arange("i", 3), j=f.new_arange("j", 4)))
assert_close(f, f(i=f.new_arange("i", 3), j="j"))
assert_close(f, f(i="i", j=f.new_arange("j", 4)))
def test_arange_2(start, stop, step):
t = randn((10, 2))
f = Tensor(t)["i"]
actual = f(i=f.new_arange("j", start, stop, step, dtype=10))
expected = Tensor(t[start:stop:step])["j"]
assert_close(actual, expected)
