fix tests

Author: cgarciae

tree_math/_src/numpy_test.py

@@ -49,7 +49,7 @@ def test_where_all_scalars(self):
     actual = tnp.where(True, 1, 2)
     self.assertTreeEqual(actual, expected, check_dtypes=False)
     with self.assertRaisesRegex(
-        TypeError, "non-tree_math.Vector argument is not a scalar",
+        TypeError, "non-tree_math.VectorBase argument is not a scalar",
     ):
       tnp.where(True, jnp.array([1, 2]), 3)
 

tree_math/_src/vector_test.py

@@ -58,7 +58,7 @@ def test_arithmetic_with_scalar(self):
     self.assertTreeEqual(vector + 1, expected, check_dtypes=True)
     self.assertTreeEqual(1 + vector, expected, check_dtypes=True)
     with self.assertRaisesRegex(
-        TypeError, "non-tree_math.Vector argument is not a scalar",
+        TypeError, "non-tree_math.VectorBase argument is not a scalar",
     ):
       vector + jnp.ones((3,))  # pylint: disable=expression-not-assigned
 
@@ -123,7 +123,7 @@ def test_matmul(self):
     self.assertAllClose(actual, expected)
 
     with self.assertRaisesRegex(
-        TypeError, "matmul arguments must both be tree_math.Vector objects",
+        TypeError, "matmul arguments must both be tree_math.VectorBase objects",
     ):
       vector1 @ jnp.ones((7,))  # pylint: disable=expression-not-assigned