Release: v0.1.0 - #10
Merged
Merged
Conversation
- SPARSE -> UNFUSED - FUSED -> FUSED_CUDA - TPMode -> MixingMode - TPNormalization -> TensorProductNormalization
test: rename CUDATensorProductParams
cleanup(core.TensorProduct): remove is_fused, is_mtpu, ... properties cleanup(TensorProduct): match / case on tp option
…e Irrep self.space iterates as (mul, ir) pairs, but the bare Irrep ir was passed directly to the Array constructor as if it were a Space, raising a TypeError on every call. Wrap it as self.space.__class__([(mul, ir)]) so each yielded block carries a proper single-block Space.
Indexing that touches the feature axis (e.g. reversing or permuting it) kept the same .space label while actually reordering irreps underneath, producing a mislabeled, non-equivariant array with no error. Raise a ValueError when an explicit index entry lands on the feature axis and isn't a full slice(None); leave plain leading-axis indexing untouched.
Neither method verified that other was the same Array subclass with a matching .space and .layout before combining arrays, silently mixing incompatible spaces/layouts (or crashing with a raw AttributeError on non-Array operands). Add a shared _check_alike guard, matching the ValueError already raised by __radd__ for the reflected case.
axis = self.feature_axis is -2 under the default TRAILING_CHANNELS layout, and other.shape[axis] raised IndexError whenever other had fewer dimensions than needed to reach that axis (e.g. a per-channel scalar of shape (channels,) against a rank-3 array). Guard the lookup with other.ndim >= -axis, treating a missing axis as an implicit broadcast size of 1 so numpy/jax's own broadcasting rules take over.
'out += " {self.layout}"' was a plain string, so it appended the
literal text "{self.layout}" instead of interpolating the actual
layout value. Add the missing f prefix.
The feature-axis guard incremented `pos` by exactly one per index-tuple element, assuming a 1:1 mapping between elements and axes. A boolean array index consumes as many axes as its own rank (e.g. a single 2-D boolean mask spans two axes), so a mask spanning the feature axis together with another axis was never checked at all, silently repackaging unrelated scalars as an equivariant array. Advance `pos` by the index's rank when it's a boolean array.
The feature-axis guard tested for Ellipsis with `in`/`.index()` and compared entries with `!= slice(None)`. For a plain numpy ndarray these fall back to elementwise `==`/`!=`, and coercing the resulting array to a bool (as `in`/`.index()`/`if` all do) raises "the truth value of an array... is ambiguous" instead of indexing correctly or raising the intended feature-axis error. jax arrays happen not to hit this since their `__eq__`/`__ne__` return NotImplemented against non-array operands, masking the bug for the common jax.Array case. Use identity/isinstance checks instead so the guard never depends on what `==`/`!=` does with the index's dtype.
_check_alike is shared by __add__ and __sub__ but always raised "Can only add arrays that are alike", so a subtraction between incompatible arrays reported the wrong operation. Pass the operation name through so the message matches the call site.
Show plain batch-axis indexing and the ValueError raised when an index would touch the feature axis.
Non contiguous padding edges (in vmap) require us to skip OOB endpoints while still looping through intermediate padding edges, but the kernel's OOB guard prevents from any time consuming work. In contrast, terminal padding edges are just dropped by the bincount and shouldn't be looped over at all. refactor: move DUMMY_INDEX to graph.py, preserve mask in vmap test: add failing vmap conv forward test fix(masking+vmap): sender case perf(convolution): drop padding edges and skip them in-kernel
`Convolution` stored `layout` as given while `TensorProduct` widens the same option to `str | Layout` and parses it, so a name reached `self.layout` unparsed and compared unequal to the enum, making the FUSED_MOSAIC_TPU guard reject a trailing-channels layout. Widen and parse to match, so every `self.layout` consumer sees a `Layout` and an unknown name fails at construction.
- add f16/f64 to the fused tp + conv kernels, drop the int32 value path - ops inherit the operand dtype: `_fused_eval` promotes x/y and packs the coefficients in it, so the buffers and the coef stride always agree - f16 not on atomicAdd paths (LEADING_CHANNELS tp, scatter_add_1): raise at trace - f64 needs jax_enable_x64, which also flips the default int dtype, so GraphCSR derives its CSR buffers as int32 explicitly - SparseMixin aggregates in the summand dtype rather than the x64 default - copy_pipe sizes by bytes now: old code over-copied at N>1 (f64 N=4 overran the buffer), f32 conv same bytes but better coalesced BREAKING CHANGE: int32 value dtype dropped from the fused kernels. scatter_add_1 with int32 values returns InvalidArgument. index dtypes (int32/uint8) unchanged.
x64 defaults sender/receiver to int64, which the handlers reject. both `_sharded_op`s cast now, so `core.Convolution` no longer has to.
we were having some impresision in on jvp on tpu. the root cause was some imprecision introduced by fact on TPU so we now use numpy for this part of the code
this change reduced drasticaly reduced the time to create Harmonics
…) (#335) * docs: add streaming message-passing animation, rendered in docs build Manim animation of the fused MP-conv forward streaming kernel (convolution/mosaic_tpu/fwd.py). It is rendered into the Sphinx site at docs build time (uvx manim step in docs.yaml -> docs/_static, gitignored) and shown on the 'Message passing on TPU' docs page so the README links to it.
a-pichard
approved these changes
Aug 14, 2026
armandpicard
approved these changes
Aug 14, 2026
olivier-peltre
force-pushed
the
release/v0.1.0
branch
from
August 14, 2026 14:11
343af62 to
eba26de
Compare
olivier-peltre
force-pushed
the
release/v0.1.0
branch
from
August 14, 2026 14:13
eba26de to
e7ce436
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See changelog