Fix projection_box crash on scalar bounds with multi-leaf pytrees

[winklemad]

No linked issue — found by reading the code.

The bug

projection_box forwards lower/upper directly into jax.tree.map(jnp.clip, tree, lower, upper), which requires them to be pytrees with the same structure as tree. So a scalar bound crashes on any multi-leaf pytree, even though the docstring documents them as "a scalar or tree with the same structure as tree":

import jax.numpy as jnp
from optax import projections as P

params = {'w': jnp.array([2.0, -1.0]), 'b': jnp.array(5.0)}

P.projection_box(params, 0.0, 1.0)   # ValueError: Expected dict, got 0.
P.projection_hypercube(params)       # ValueError: Expected dict, got 0.  (default scale=1)

projection_hypercube's default scale=1 is a scalar, so projection_hypercube(params) fails on the common case of a multi-leaf parameter tree (single-leaf inputs happen to work, which is why the existing tests never caught it).

Smoking gun

The sibling projection_linf_ball already works around this by broadcasting its scalar scale to the tree before calling projection_box:

lower = optax.tree.full_like(tree, -scale)
upper = optax.tree.full_like(tree, scale)
return projection_box(tree, lower=lower, upper=upper)

so projection_linf_ball(params, 1.0) succeeds on the exact multi-leaf tree that projection_hypercube(params) crashes on.

Fix

Broadcast a scalar bound to the tree structure inside projection_box (a bound that already matches the tree's structure is passed through unchanged, so both a scalar and a matching tree are accepted). This fixes projection_box for scalar bounds and, transitively, projection_hypercube.

Testing

Added sub-tests for scalar bounds on a multi-leaf tree to both test_projection_box and test_projection_hypercube. They fail before the change (ValueError) and pass after; the full optax/projections/_projections_test.py suite passes (39 tests / 43 sub-tests). ruff is clean and the changed lines are pyink-formatted.