Hey! Continuing the discussion about VectorMixin for pytree classes, I see two ways to implement this:

1. VectorMixin inherits from Vector and overrides the tree property to return self. Roughly implemented as:

class VectorMixin(tm.Vector):
    """A mixin class for vector operations that works with any pytree class"""

    @property
    def tree(self):
        return self

2. VectorMixin does not inherit from Vector, instead of using .tree it assumes self is a pytree and implements all the operators based on this. Vector inherits from VectorMixin and implements the pytree protocol roughly as:

@tree_util.register_pytree_node_class
class Vector(VectorMixin):
  """A wrapper for treating an arbitrary pytree as a 1D vector."""

  def __init__(self, tree):
    self._tree = tree

  def tree_flatten(self):
    return (self._tree,), None

  @classmethod
  def tree_unflatten(cls, _, args):
    return cls(*args)
  
  ...

First option is the easiest, second requires a refactor but its more in the spirit of what a mixin "should be". WDYT?

I'm not sure if I'm using this right, but here's what I have:

from tree_math import Vector
from typing import Any, Callable, TypeVar
from jax.scipy.optimize import minimize

T = TypeVar('T')


def tree_minimize(fun: Callable[[T], RealNumeric], x0: T,
                  *,
                  method: str,
                  tol: float | None = None,
                  options: None | Mapping[str, Any] = None) -> T:
    wrapped = Vector(x0)
    def wrapped_fun(x: Vector) -> RealNumeric:
        return fun(x.tree)
    optimize_result = minimize(wrapped_fun, wrapped, method=method, tol=tol, options=options)
    return optimize_result.x.tree  # type: ignore[attr-defined]

When I call this, I get:

  File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.10/lib/python3.10/site-packages/jax/_src/scipy/optimize/minimize.py", line 103, in minimize
    results = minimize_bfgs(fun_with_args, x0, **options)
  File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.10/lib/python3.10/site-packages/jax/_src/scipy/optimize/bfgs.py", line 102, in minimize_bfgs
    converged=jnp.linalg.norm(g_0, ord=norm) < gtol,
  File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.10/lib/python3.10/site-packages/jax/_src/numpy/linalg.py", line 438, in norm
    x, = _promote_dtypes_inexact(jnp.asarray(x))
  File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.10/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1924, in asarray
    return array(a, dtype=dtype, copy=False, order=order)
  File "/home/neil/.cache/pypoetry/virtualenvs/cmm-tspD8tmv-py3.10/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1903, in array
    raise TypeError(f"Unexpected input type for array: {type(object)}")
TypeError: Unexpected input type for array: <class 'tree_math._src.vector.Vector'>

I thought Vector was supposed to pretend to be a jax.Array? Is it impossible to add a __jax_array__ method to tree_math.VectorMixin?

As discussed in #3 this PR adds the VectorBase class which both contains the base vector logic and can act as a mixin that can be applied to any pytree class. The Vector class is now a simple pytree container that inherits from VectorBase.

I've use VectorBase here as it reads better but we can still consider the VectorMixin option.

Pending:

• [x] Add a test to check that this works with classes other than Vector.

Refactor navier_stokes_explicit_terms() out of semi_implicit_navier_stokes().

Enables additional options for writing systems of equations that use explicit Navier-Stokes, i.e. by writing a custom state via tree_math.

import jax.numpy as jnp 
import tree_math as tm

def f(x, y):
  return x, y
  
x = y = tm.Vector(jnp.array(0.))

tm.unwrap(f, out_vectors = (True, False))(x, y)
# (tree_math.Vector(DeviceArray(0., dtype=float32, weak_type=True)), DeviceArray(0., dtype=float32, weak_type=True))
tm.unwrap(f, out_vectors = [True, False])(x, y)
# ValueError: Expected list, got (DeviceArray(0., dtype=float32, weak_type=True), DeviceArray(0., dtype=float32, weak_type=True)).

How should I go about importing struct? Thanks!

My attempt below fails --

import tree_math

@tree_math.struct
class Point: 
    x: float 
    y: float 

AttributeError: module 'tree_math' has no attribute 'struct'

Wondering how well tree-math supports computation on multiple devices?

Let's say we have a pytree of tensors of different dimensions and want to perform some operations on each of them with tree-math, can we distribute those tasks to multiple devices (GPU, for instance)?

It would be nice to have an easy way to define dataclasses that are also tree-math vectors.

We could borrow the syntax of flax.struct here: https://flax.readthedocs.io/en/latest/flax.struct.html

Example usage:

from tree_math import struct

@struct
class FluidState:
  velocity_x: Array
  velocity_y: Array
  pressure: Array

It seems like the current implementation doesn't allow broadcasting arguments. Here's an example for normalizing leafs.

import tree_math as tm
import jax
import jax.numpy as jnp

a = jnp.ones(10)
b = jnp.ones(5)

v = tm.Vector({'a': a, 'b': b})
v / jax.tree_map(jnp.linalg.norm, v)

returns the following error:

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-19-c2a8ad9c2f8f> in <module>()
----> 1 v / jax.tree_map(jnp.linalg.norm, v)

2 frames
/usr/local/lib/python3.7/dist-packages/tree_math/_src/vector.py in wrapper(self, other)
     72   """Implement a forward binary method, e.g., __add__."""
     73   def wrapper(self, other):
---> 74     return broadcasting_map(func, self, other)
     75   wrapper.__name__ = f"__{name}__"
     76   return wrapper

/usr/local/lib/python3.7/dist-packages/tree_math/_src/vector.py in broadcasting_map(func, *args)
     65   if not vector_args:
     66     return func2()  # result is a scalar
---> 67   _flatten_together(*[arg.tree for arg in vector_args])  # check shapes
     68   return tree_util.tree_map(func2, *vector_args)
     69 

/usr/local/lib/python3.7/dist-packages/tree_math/_src/vector.py in _flatten_together(*args)
     37   if not all(shapes == all_shapes[0] for shapes in all_shapes[1:]):
     38     shapes_str = " vs ".join(map(str, all_shapes))
---> 39     raise ValueError(f"tree leaves have different array shapes: {shapes_str}")
     40 
     41   return all_values, all_treedefs[0]

ValueError: tree leaves have different array shapes: [(10,), (5,)] vs [(), ()]