mirror of
https://github.com/firestar5683/StarPilot.git
synced 2026-11-12 15:33:45 +08:00
firestar5683
22a8575199
Keep official openpilot ancestry and dependency licenses. Store resources directly in Git, remove LFS setup requirements, and leave the downloadable big model outside the source tree.
172 lines
6.1 KiB
Python
172 lines
6.1 KiB
Python
# modified from
|
|
# https://github.com/arogozhnikov/einops/blob/master/tests/test_examples.py
|
|
# https://github.com/arogozhnikov/einops/blob/master/tests/test_ops.py
|
|
# https://github.com/arogozhnikov/einops/blob/master/tests/test_parsing.py
|
|
|
|
import numpy as np
|
|
import unittest
|
|
from tinygrad import Tensor
|
|
|
|
|
|
class test_rearrange_examples(unittest.TestCase):
|
|
def test1(self):
|
|
# transpose
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
y = x.rearrange("b c h w -> b h w c")
|
|
assert tuple(y.shape) == (10, 30, 40, 20)
|
|
|
|
def test2(self):
|
|
# view / reshape
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
y = x.rearrange("b c h w -> b (c h w)")
|
|
assert tuple(y.shape) == (10, 20 * 30 * 40)
|
|
|
|
def test3(self):
|
|
# depth-to-space
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
y = x.rearrange("b (c h1 w1) h w -> b c (h h1) (w w1)", h1=2, w1=2)
|
|
assert tuple(y.shape) == (10, 5, 30 * 2, 40 * 2)
|
|
|
|
def test4(self):
|
|
# space-to-depth
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
y = x.rearrange("b c (h h1) (w w1) -> b (h1 w1 c) h w", h1=2, w1=2)
|
|
assert tuple(y.shape) == (10, 20 * 4, 30 // 2, 40 // 2)
|
|
|
|
def test5(self):
|
|
# simple transposition
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
y = x.rearrange("b1 sound b2 letter -> b1 b2 sound letter")
|
|
assert tuple(y.shape) == (10, 30, 20, 40)
|
|
|
|
def test6(self):
|
|
# parsing parameters
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
t = x.rearrange("b c h w -> (b h w) c")
|
|
t = t[:, ::2] # replacement for dot-product, just changes size of second axis
|
|
assert tuple(t.shape) == (10 * 30 * 40, 10)
|
|
|
|
def test7(self):
|
|
x = Tensor(np.arange(10 * 20 * 30 * 40, dtype=np.int32).reshape([10, 20, 30, 40]))
|
|
# split of embedding into groups
|
|
y1, y2 = x.rearrange("b (c g) h w -> g b c h w", g=2)
|
|
assert tuple(y1.shape) == (10, 10, 30, 40)
|
|
assert tuple(y2.shape) == (10, 10, 30, 40)
|
|
|
|
def test8(self):
|
|
x = Tensor(np.arange(10 * 20 * 1 * 1, dtype=np.int32).reshape([10, 20, 1, 1]))
|
|
# squeeze - unsqueeze
|
|
y = x.rearrange("b c () () -> b c")
|
|
assert tuple(y.shape) == (10, 20)
|
|
y = y.rearrange("b c -> c b () ()")
|
|
assert tuple(y.shape) == (20, 10, 1, 1)
|
|
|
|
def test9(self):
|
|
x = Tensor(np.arange(10 * 20 * 1 * 1, dtype=np.int32).reshape([10, 20, 1, 1]))
|
|
# squeeze - unsqueeze
|
|
y = x.rearrange("b c 1 1 -> b c")
|
|
assert tuple(y.shape) == (10, 20)
|
|
y = y.rearrange("b1 c -> c b1 1 1")
|
|
assert tuple(y.shape) == (20, 10, 1, 1)
|
|
|
|
|
|
class test_rearrange_ops(unittest.TestCase):
|
|
def test_rearrange_errors(self):
|
|
x = Tensor.zeros([1, 1, 1, 1, 1])
|
|
x.rearrange("a b c d ... -> a b c ... d")
|
|
bad_patterns = [
|
|
"a b c d (...) -> a b c ... d", # collapsed ellipsis on input
|
|
"a b (c d ... -> a b c ... d", # unbalanced brackets
|
|
"a b* c d ... -> a b c ... d", # not alphanumeric
|
|
"a b c d -> a b c d -> a b c d", # two "->"
|
|
"a ... c ... -> ... a ... c", # two "..."
|
|
"a b c d e -> f b c d e", # name mismatch
|
|
]
|
|
for pattern in bad_patterns:
|
|
with self.assertRaises(AssertionError):
|
|
x.rearrange(pattern)
|
|
|
|
x.rearrange("... -> (...)")
|
|
with self.assertRaises(AssertionError):
|
|
x.rearrange("(...) -> (...)")
|
|
|
|
y = Tensor.zeros([8, 1])
|
|
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2, a2=2)
|
|
with self.assertRaises(RuntimeError):
|
|
## should fail as not enough dimensions specified
|
|
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2)
|
|
with self.assertRaises(ValueError):
|
|
## should fail as 6 does not divide 8
|
|
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=3, a2=2)
|
|
with self.assertRaises(AssertionError):
|
|
## incorrect dimension provided for an axis that is only permuted
|
|
y.rearrange("(a1 a2 a3) b -> b a3 a2 a1", a1=2, a2=2, b=2)
|
|
with self.assertRaises(AssertionError):
|
|
## unused axis provided
|
|
y.rearrange("(a b c) d -> a b c d", b=2, c=2, e=2)
|
|
|
|
|
|
class test_rearrange_parsing(unittest.TestCase):
|
|
def test_elementary_axis_name(self):
|
|
for name in [
|
|
"a",
|
|
"b",
|
|
"h",
|
|
"dx",
|
|
"h1",
|
|
"zz",
|
|
"i9123",
|
|
"somelongname",
|
|
"Alex",
|
|
"camelCase",
|
|
"u_n_d_e_r_score",
|
|
"unreasonablyLongAxisName",
|
|
]:
|
|
Tensor.ones((1,)).rearrange(f"{name} -> {name}")
|
|
|
|
for name in ["2b", "12", "_startWithUnderscore", "endWithUnderscore_", "_"]:
|
|
with self.assertRaises(AssertionError):
|
|
Tensor.ones((1,)).rearrange(f"{name} -> {name}")
|
|
|
|
with self.assertRaises(RuntimeError):
|
|
Tensor.ones((1,)).rearrange(" -> ")
|
|
|
|
def test_invalid_expressions(self):
|
|
# double ellipsis should raise an error
|
|
def _test_expression(expression: str):
|
|
Tensor.ones((2, 3, 4, 5, 6)).rearrange(f"{expression} -> {expression}")
|
|
|
|
_test_expression("... a b c d")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("... a b c d ...")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("... a b c (d ...)")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("(... a) b c (d ...)")
|
|
|
|
# double/missing/enclosed parenthesis
|
|
Tensor.ones((2, 3, 4, 5, 6)).rearrange("a b c d ... -> (a) b c (d ...)")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("(a)) b c (d ...)")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("(a b c (d ...)")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("(a) (()) b c (d ...)")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("(a) ((b c) (d ...))")
|
|
|
|
# invalid identifiers
|
|
_test_expression("camelCase under_scored cApiTaLs ß ...")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("1a")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("_pre")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("...pre")
|
|
with self.assertRaises(AssertionError):
|
|
_test_expression("pre...")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|