fix: restore jax.tree_util.Partial in omega() for PowerLaw+JAX#352

Merged

Jammy2211 merged 1 commit intomainfrom

feature/group-features

Apr 14, 2026

Commits on Apr 14, 2026

fix: restore jax.tree_util.Partial in omega() for PowerLaw+JAX

authored and
Jammy2211
committed