Fix deferred rendering for PRNG key arrays

[abhinav-bellapu]

Summary

• render typed PRNG key arrays through their numeric key data before truncation
• preserve JAX's native typed-key representation without materializing large arrays in full
• cover every built-in PRNG implementation for one- and two-dimensional key batches

Root cause

The deferred representation path passed JAX's extended PRNG key dtype directly to NumPy. NumPy cannot represent that dtype, so expanding a sufficiently large key array raised TypeError instead of rendering it.

Fixes #71.

Validation

• uv run pyink treescope tests --check
• uv run pylint treescope
• uv run python run_tests.py (180 passed, 1 skipped)
• uv run pytype --jobs auto treescope

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.