refactor into VectorBase

Author: cgarciae

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, Vector)]
+  static_argnums = [i for i, x in enumerate(args) if not isinstance(x, VectorBase)]
   func2, vector_args = _argnums_partial(func, args, static_argnums)
   for arg in args:
-    if not isinstance(arg, Vector):
+    if not isinstance(arg, VectorBase):
       shape = jnp.shape(arg)
       if shape:
         raise TypeError(
-            f"non-tree_math.Vector argument is not a scalar: {arg!r}"
+            f"non-tree_math.VectorBase 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, Vector) or not isinstance(right, Vector):
-    raise TypeError("matmul arguments must both be tree_math.Vector objects")
+  if not isinstance(left, VectorBase) or not isinstance(right, VectorBase):
+    raise TypeError("matmul arguments must both be tree_math.VectorBase objects")
 
   def _vector_dot(a, b):
     return jnp.dot(jnp.ravel(a), jnp.ravel(b), precision=precision)
@@ -206,7 +206,7 @@ def max(self):
     return jnp.asarray(list(parts)).max()
 
 @tree_util.register_pytree_node_class
-class Vector:
+class Vector(VectorBase):
   """A wrapper for treating an arbitrary pytree as a 1D vector."""
 
   def __init__(self, tree):