assert to prepare for grad uop [pr] (#8280)

* assert to prepare for grad uop [pr]

* fix test_nn

* fix most of test_tensor

* few more tests

* fix multi

* uniform gradient

* acc_dtype

* any for multi

* fix typing

* fix assert, CAST_BEFORE_VIEW is still the issue

* explict test for CAST_BEFORE_VIEW

---------

Co-authored-by: qazal <77887910+Qazalin@users.noreply.github.com>
This commit is contained in:
George Hotz
2025-01-14 13:26:56 -08:00
committed by GitHub
parent fdd46c9f28
commit c85737c200
5 changed files with 20 additions and 2 deletions

View File

@@ -781,7 +781,8 @@ class TestAutoCastType(unittest.TestCase):
if DEBUG >= 2:
print(f"testing {default_dtype=}, {dtype=}")
a = Tensor([1, 2, 3], dtype=dtype, requires_grad=True)
b = (a * 5).sum()
# NOTE: this is broken without default_dtype because of CAST_BEFORE_VIEW
b = (a * 5).sum(acc_dtype=default_dtype)
b.backward() # if there is dtype mismatch, lazy should assert
assert a.grad.dtype == a.dtype
np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])