rename VectorBase to VectorMixin

Author: cgarciae

tree_math/__init__.py

@@ -20,7 +20,7 @@
     wrap,
     unwrap,
 )
-from tree_math._src.vector import Vector, VectorBase
+from tree_math._src.vector import Vector, VectorMixin
 import tree_math.numpy
 
 __version__ = '0.1.0'

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.VectorBase argument is not a scalar",
+        TypeError, "non-tree_math.VectorMixin argument is not a scalar",
     ):
       tnp.where(True, jnp.array([1, 2]), 3)
 

tree_math/_src/vector.py

@@ -53,14 +53,14 @@ def g(*args3):
 
 def broadcasting_map(func, *args):
   """Like tree_map, but scalar arguments are broadcast to all leaves."""
-  static_argnums = [i for i, x in enumerate(args) if not isinstance(x, VectorBase)]
+  static_argnums = [i for i, x in enumerate(args) if not isinstance(x, VectorMixin)]
   func2, vector_args = _argnums_partial(func, args, static_argnums)
   for arg in args:
-    if not isinstance(arg, VectorBase):
+    if not isinstance(arg, VectorMixin):
       shape = jnp.shape(arg)
       if shape:
         raise TypeError(
-            f"non-tree_math.VectorBase argument is not a scalar: {arg!r}"
+            f"non-tree_math.VectorMixin argument is not a scalar: {arg!r}"
         )
   if not vector_args:
     return func2()  # result is a scalar
@@ -112,8 +112,8 @@ def dot(left, right, *, precision="highest"):
   Returns:
     Resulting dot product (scalar).
   """
-  if not isinstance(left, VectorBase) or not isinstance(right, VectorBase):
-    raise TypeError("matmul arguments must both be tree_math.VectorBase objects")
+  if not isinstance(left, VectorMixin) or not isinstance(right, VectorMixin):
+    raise TypeError("matmul arguments must both be tree_math.VectorMixin objects")
 
   def _vector_dot(a, b):
     return jnp.dot(jnp.ravel(a), jnp.ravel(b), precision=precision)
@@ -122,7 +122,7 @@ def _vector_dot(a, b):
   parts = map(_vector_dot, left_values, right_values)
   return functools.reduce(operator.add, parts)
 
-class VectorBase:
+class VectorMixin:
   """A mixin class that adds a 1D vector-like behaviour to any custom pytree class."""
 
   @property
@@ -206,7 +206,7 @@ def max(self):
     return jnp.asarray(list(parts)).max()
 
 @tree_util.register_pytree_node_class
-class Vector(VectorBase):
+class Vector(VectorMixin):
   """A wrapper for treating an arbitrary pytree as a 1D vector."""
 
   def __init__(self, tree):

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.VectorBase argument is not a scalar",
+        TypeError, "non-tree_math.VectorMixin 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.VectorBase objects",
+        TypeError, "matmul arguments must both be tree_math.VectorMixin objects",
     ):
       vector1 @ jnp.ones((7,))  # pylint: disable=expression-not-assigned
 
@@ -151,7 +151,7 @@ def test_sum_mean_min_max(self):
   def test_custom_class(self):
 
     @tree_util.register_pytree_node_class
-    class CustomVector(tm.VectorBase):
+    class CustomVector(tm.VectorMixin):
 
       def __init__(self, a: int, b: float):
         self.a = a