ITADN

Bug: treescope fails to render typed PRNG key arrays

#52Closedamifalk 创建于 2025-01-03
bug
A
amifalkcommented
With jax 0.4.37 and treescope 0.1.7: ```python import jax import treescope as ts ts.basic_interactive_setup() jax.random.key(0) ``` ``` <TypeError during deferred rendering Traceback (most recent call last): File ".../.venv/lib/python3.10/site-packages/treescope/lowering.py", line 358, in _render_to_html_as_root_streaming replacement_part = deferred.thunk(layout_decision) File ".../.venv/lib/python3.10/site-packages/treescope/external/jax_support.py", line 606, in _thunk summarized = adapter.get_array_summary(node, fast=False) File ".../.venv/lib/python3.10/site-packages/treescope/external/jax_support.py", line 525, in get_array_summary output_parts.append(summarize_array_data(array)) File ".../.venv/lib/python3.10/site-packages/treescope/external/jax_support.py", line 479, in summarize_array_data output_parts.extend(_summarize_array_data_unconditionally(array)) File ".../.venv/lib/python3.10/site-packages/treescope/external/jax_support.py", line 433, in _summarize_array_data_unconditionally stat = compute_summary(array, is_floating, is_integer, is_bool) File ".../.venv/lib/python3.10/site-packages/treescope/external/jax_support.py", line 391, in _compute_summary x = xnp.array(x) File ".../.venv/lib/python3.10/site-packages/jax/_src/prng.py", line 283, 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. > ```
关闭于 2025-01-21 1 条评论