tree-math

library0.2.1pypypi✓ verified 85d ago

Provides mathematical operations (add, multiply, etc.) on JAX pytrees, treating them as vectors. Version 0.2.1, infrequent releases, currently active.

ai-ml

pip install tree-math

github ↗package ↗

INSTALL

IMPORT

SIG · TREE-MATH

tree-math

ai-mlpythonv0.2.1

harness data pending

Install & Compatibility

Where this runs

No compatibility data collected yet for this library.

Code

Verified usage

Verified import paths — ran on the pinned version, not inferred.

tree_math

✓ import tree_math

✗ from tree_math import *

Star imports can cause name collisions; prefer explicit import of required classes.

Vector

✓ from tree_math import Vector

✗ import tree_math.Vector

Vector is a class, not a submodule.

Creates a Vector from a pytree of arrays and performs basic operations.

import jax.numpy as jnp from tree_math import Vector tree = {'a': jnp.array([1.0, 2.0]), 'b': jnp.array([3.0])} v = Vector(tree) print(v + v) print(2 * v) print(v @ v)

Debug

Known issues

gotchaAll leaf arrays must have the same dtype and be convertible to a single type. Mixing float32 and float64 can cause unexpected type promotion errors.

fix

Ensure all arrays in the pytree have consistent dtypes before constructing a Vector.

affects: all

gotchaVector.__init__ does not perform a deep copy of the pytree; mutations to the original tree may affect the Vector.

fix

Use jax.tree_util.tree_map(lambda x: x.copy(), tree) if you need independent copies.

affects: all

Errors

Common errors & fixes

ModuleNotFoundError: No module named 'tree_math'

tree-math package not installed.

fix

Run 'pip install tree-math' in your virtual environment.

ValueError: All leaves must have the same dtype.

The pytree contains arrays with different dtypes (e.g., float32 and float64).

fix

Convert all arrays to the same dtype, e.g., tree = jax.tree_util.tree_map(lambda x: x.astype(jnp.float32), tree).

Upgrade

Version history

0.2.1latest on PyPI · released Apr 4, 2024

Audit

Dependencies

jaxrequiredcore dependency for pytree support

Agent activity

11 hits · last 30 days

node

OpenAI (training)

Resources

githubhttps://github.com/google/tree-math ↗

packagehttps://pypi.org/project/tree-math/ ↗