Make rkg(1) return the same shape as rkg(n)

[cemde]

Bug

def __call__(self, n: int = 1):
    _k = self.key.split(n)
    if n == 1:
        return _k[0]      # unwraps to a bare key
    return _k

The convenience unwrap for a single key was keyed off n == 1, which is also the default. So "give me one key" and "give me a batch of n keys" were conflated at n == 1.

Impact

rkg(n) returns shape (n, 2) for n >= 2 but a bare key of shape (2,) for n == 1. So keys[0] is a uint32 rather than a key, and iterating yields the key's two integers rather than one key. Any batched path hits this on a final partial batch of size 1, or a single-example debug run, and has to special-case it with nothing in the API to suggest so.

Fix

Make the default a None sentinel, so the unwrap keys off "no argument given" rather than off the value 1:

def __call__(self, n: int | None = None):
    _k = self.key.split(1 if n is None else n)
    if n is None:
        return _k[0]
    return _k

rkg() is unchanged. rkg(1) now returns (1, 2) like every other n.

Why the signature had to change

The two contracts are only separable if the default stops being the literal 1. test_keys_are_usable_by_jax_random does jax.random.normal(rkg(), (16,)), and jax rejects a (1, 2) key there, so the no-argument call must keep returning a bare key.

Breaking change

rkg(1) previously returned a scalar-indexable bare key, and one shipped example relied on that. examples/7_mcpc.ipynb had seed: int = lambda: px.RKG(1)[0], where [0] on a (2,) key produced the scalar seed optax.add_noise wants. Under the new contract that expression yields shape (2,) and optax raises TypeError: key accepts a scalar seed.

Updated in this PR to px.RKG()[0], which is the no-argument form and gives the scalar back. Any downstream code calling rkg(1) to mean "one bare key" needs the same edit.

Both internal callers (Linear and Conv2d, via pcx/nn/_layer.py) use the no-argument form and are unaffected. Vmap._t manipulates RKG.key.split(...) directly and bypasses __call__.

Closes #75

[liukidar]

Lgtm