v0.5.0

Highlights

jaxmore.loop: one call, several ways to loop

jaxmore.loop(strategy, f, signature=...) maps a function over the broadcast
batch dimensions of its arguments. The strategy, from jaxmore.loop_strategies,
decides how the loop runs:

Strategy	Runs as
Vectorize	one jax.vmap per batch dimension (like jnp.vectorize)
VMap	same as Vectorize; a separate tag so libraries can dispatch on it
Scan	jax.lax.map per batch dimension: sequential, one call's memory
NoLoop	f unchanged, for an already-batched function
Determine	Scan on CPU, Vectorize elsewhere, decided when called

import jax.numpy as jnp
import jaxmore as jm
from jaxmore import loop_strategies as ls

norm_t = lambda x, t: jnp.linalg.norm(x) * t
jm.loop(ls.Scan, norm_t, signature="(n),()->()")(jnp.ones((4, 3)), jnp.arange(4.0))

• signature is a NumPy gufunc signature. Named core dimensions must agree
  across arguments, and numbers fix a size.
• Positional arguments may be pytrees. Keyword arguments are passed to f
  unmapped. Outputs are inferred from f and may be any pytree.
• No strategy broadcasts an argument, so one shared by the whole batch is never
  copied.

💥 Breaking changes

• JAX >= 0.7.2 and Python >= 3.11 are now required (#36). JAX 0.7 dropped
  Python 3.10, so the two floors go together. If you are on Python 3.10 or an
  older JAX, stay on jaxmore<0.5.

What's Changed

• ✨ feat(loop): add loop and loop_strategies by @nstarman in #35
• ⬆️ dep-bump: require jax>=0.7.2 and Python>=3.11 by @nstarman in #36
• 👷 ci: don't cancel CI runs on main by @nstarman in #34
• 🎨 assets: add a logo, JAX's with the J's stem crossed into a + by @nstarman in #28
• ci(prek): protect main and versions/ branches from direct commits by @nstarman in #26
• 💚 ci: bump hynek/build-and-inspect-python-package to v3.0.1 by @nstarman in #25

The bounded_while_loop performance fix (#30) shipped in v0.4.1 and is not new
here.

Full Changelog: v0.4.1...v0.5.0