Accept the argnums and switch spellings the signatures promise

[cemde]

Bug

Two places where the annotation is wider than what the code accepts.

ValueAndGrad is typed argnums: int | Sequence[int], matching jax, but _t splats it: jax.value_and_grad(..., argnums=(*argnums, n)). An int is not iterable.

Switch is typed Sequence[...] and passes fns to _BaseTransform.__init__, which calls _make_tuple(fn). That helper is x if isinstance(x, tuple) else (x,), so a list of branches is wrapped rather than expanded, collapsing all branches into one element.

Impact

Both documented spellings raise:

pxf.value_and_grad(argnums=0)     TypeError: Value after * must be an iterable, not int
pxf.switch([b0, b1, b2])          TypeError: object of type 'function' has no len()

Only the undocumented argnums=(0,) and a tuple of branches worked.

Fix

Normalise at each call site: ValueAndGrad.__init__ wraps an int and tuples anything else, and Switch.__init__ calls super().__init__(tuple(fns)).

Normalising in __init__ rather than _t also fixes a silent bug: a generator argnums worked on the first call and then returned no positional gradient at all on the second, because _t re-splatted the now-exhausted generator.

Why not in _make_tuple

Teaching the shared helper to expand sequences breaks four of its six callers. The decisive one is _process_mask, which passes mask dict keys: str is a Sequence, so "model" would expand into five single-character keys and the mask would fail with a pytree structure error. Two others pass user function return values, where a returned list would be resplit into value plus aux.

Closes #74

[liukidar]

LGTM