PRNGKey dtype cannot be converted to a NumPy array

I have a NNX module that is attempted to be used as an array. The module has an nnx.rngs has an instance variable on the class which is causing the following. ``` /usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/renderers.py:314: UserWarning: Ignoring error while formatting value of type <class 'jax._src.prng.PRNGKeyArray'> with <function handle_via_global_registry at 0x320e37837d80>: Traceback (most recent call last): File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/renderers.py", line 290, in _render_subtree maybe_result = handler(node=node, path=path, subtree_renderer=rec) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/_internal/handlers/custom_type_handlers.py", line 100, in handle_via_global_registry return maybe_handler(node, path, subtree_renderer) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/external/jax_support.py", line 610, in render_jax_arrays renderable=lowering.maybe_defer_rendering( ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/lowering.py", line 91, in maybe_defer_rendering return main_thunk(None) ^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/external/jax_support.py", line 582, in _thunk summarized = adapter.get_array_summary(node, fast=False) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/external/jax_support.py", line 501, in get_array_summary output_parts.append(summarize_array_data(array)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/external/jax_support.py", line 455, in summarize_array_data output_parts.extend(_summarize_array_data_unconditionally(array)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/external/jax_support.py", line 409, in _summarize_array_data_unconditionally stat = compute_summary(array, is_floating, is_integer, is_bool) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/treescope/external/jax_support.py", line 367, in _compute_summary x = xnp.array(x) ^^^^^^^^^^^^ File "/usr/local/google/_blaze_joetoth/4e3a01cc332ff2bd5735d76fbcbe09fc/execroot/google3/blaze-out/haswell-opt/bin/learning/autolx/ingestables/core/v3/examples/ingestables/main.runfiles/google3/third_party/py/jax/_src/prng.py", line 282, in __array__ raise TypeError("JAX array with PRNGKey dtype cannot be converted to a NumPy array." TypeError: JAX array with PRNGKey dtype cannot be converted to a NumPy array. Use jax.random.key_data(arr) if you wish to extract the underlying integer array. ```