support custom UOp kernels (#13028)

* support custom UOp kernels

* no number

* multioutput works

* backward kernel runs

* move kernel class

* grad later

* work

* no tags in kernel graph

* test arange

* arange + contig

* delete comment
This commit is contained in:
George Hotz
2025-10-31 15:51:39 +08:00
committed by GitHub
parent 9f0c25ec48
commit b791d70725
6 changed files with 107 additions and 24 deletions

View File

@@ -711,7 +711,7 @@ class TestSchedule(unittest.TestCase):
self.assertEqual(b.buffer.numpy(), [12])
# unlike schedule, kernelize can be called multiple times on a Tensor
def test_double_kerenlize(self):
def test_double_kernelize(self):
a = Tensor.empty(10)
b = Tensor.empty(10)
c = (a+b)
@@ -2267,7 +2267,7 @@ class TestContiguous(unittest.TestCase):
def test_double_contiguous_realizes_once(self):
a = Tensor.empty(4, 1)
b = a.expand((4, 4)).contiguous().contiguous()
check_schedule(b, 2) # TODO: should be 1?
check_schedule(b, 1)
def test_view_does_not_realize(self):
a = Tensor.empty(4)