def test_match_equation_dtype(self): # A simple coercion eqns = _match_equation(dshape('int32'), dshape('int64')) self.assertEqual(eqns, [(T.int32, T.int64)]) # Matching a data type variable eqns = _match_equation(dshape('int32'), dshape('D')) self.assertEqual(eqns, [(T.int32, T.TypeVar('D'))])
def test_match_equation_dim(self): # Broadcasting a single dimension eqns = _match_equation(dshape('1 * int32'), dshape('10 * int32')) self.assertEqual(eqns, [(T.Fixed(1), T.Fixed(10)), (T.int32, T.int32)]) # Matching a dim type variable eqns = _match_equation(dshape('3 * int32'), dshape('M * int32')) self.assertEqual(eqns, [(T.Fixed(3), T.TypeVar('M')), (T.int32, T.int32)])
def test_match_equation_ellipsis(self): # Matching an ellipsis eqns = _match_equation(dshape('int32'), dshape('... * int32')) self.assertEqual(eqns, [([], T.Ellipsis()), (T.int32, T.int32)]) eqns = _match_equation(dshape('3 * int32'), dshape('... * int32')) self.assertEqual(eqns, [([T.Fixed(3)], T.Ellipsis()), (T.int32, T.int32)]) eqns = _match_equation(dshape('3 * var * int32'), dshape('... * int32')) self.assertEqual(eqns, [([T.Fixed(3), T.Var()], T.Ellipsis()), (T.int32, T.int32)]) # Matching an ellipsis type variable eqns = _match_equation(dshape('int32'), dshape('A... * int32')) self.assertEqual(eqns, [([], T.Ellipsis(T.TypeVar('A'))), (T.int32, T.int32)]) eqns = _match_equation(dshape('3 * int32'), dshape('A... * int32')) self.assertEqual(eqns, [([T.Fixed(3)], T.Ellipsis(T.TypeVar('A'))), (T.int32, T.int32)]) eqns = _match_equation(dshape('3 * var * int32'), dshape('A... * int32')) self.assertEqual(eqns, [([T.Fixed(3), T.Var()], T.Ellipsis(T.TypeVar('A'))), (T.int32, T.int32)]) # Matching an ellipsis with a dim type variable on the left eqns = _match_equation(dshape('3 * var * int32'), dshape('A * B... * int32')) self.assertEqual(eqns, [(T.Fixed(3), T.TypeVar('A')), ([T.Var()], T.Ellipsis(T.TypeVar('B'))), (T.int32, T.int32)]) # Matching an ellipsis with a dim type variable on the right eqns = _match_equation(dshape('3 * var * int32'), dshape('A... * B * int32')) self.assertEqual(eqns, [([T.Fixed(3)], T.Ellipsis(T.TypeVar('A'))), (T.Var(), T.TypeVar('B')), (T.int32, T.int32)]) # Matching an ellipsis with a dim type variable on both sides eqns = _match_equation(dshape('3 * var * int32'), dshape('A * B... * C * int32')) self.assertEqual(eqns, [(T.Fixed(3), T.TypeVar('A')), ([], T.Ellipsis(T.TypeVar('B'))), (T.Var(), T.TypeVar('C')), (T.int32, T.int32)]) eqns = _match_equation(dshape('3 * var * 4 * M * int32'), dshape('A * B... * C * int32')) self.assertEqual(eqns, [(T.Fixed(3), T.TypeVar('A')), ([T.Var(), T.Fixed(4)], T.Ellipsis(T.TypeVar('B'))), (T.TypeVar('M'), T.TypeVar('C')), (T.int32, T.int32)])