mirror of
https://jihulab.com/mr-one/openpilot.git
synced 2026-10-28 21:43:44 +08:00
f8e13417f2
date: 2026-08-15T14:20:10 master commit: d48cfa730bf21593b684efeaf3fd8940695e1dea
30 lines
750 B
Python
30 lines
750 B
Python
import unittest
|
|
from tinygrad.tensor import Tensor
|
|
|
|
class TestMaskedTensor(unittest.TestCase):
|
|
def test_mul_masked(self):
|
|
a = Tensor([1,1,1,1,1])
|
|
b = Tensor([1,1]).pad(((0,3),))
|
|
c = a*b
|
|
assert c.shape == a.shape
|
|
ret = c.data()
|
|
assert ret.tolist() == [1.0, 1.0, 0.0, 0.0, 0.0]
|
|
|
|
def test_mul_both_masked(self):
|
|
a = Tensor([1,1]).pad(((0,3),))
|
|
b = Tensor([1,1]).pad(((0,3),))
|
|
c = a*b
|
|
assert c.shape == a.shape
|
|
ret = c.data()
|
|
assert ret.tolist() == [1.0, 1.0, 0.0, 0.0, 0.0]
|
|
|
|
def test_add_masked(self):
|
|
a = Tensor([1,1]).pad(((0,2),))
|
|
b = Tensor([1,1]).pad(((0,2),))
|
|
c = a+b
|
|
ret = c.data()
|
|
assert ret.tolist() == [2.0, 2.0, 0.0, 0.0]
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|