FR: Dimension dependent axis-item labels for treescope.render_array

[amifalk]

Currently, the axis-item labels for an array cannot depend on multiple dimensions, but sometimes the meaning of an axis depends on another dimension that is sliced into.

This interpretation should be coherent so long as the axis-item labels named array is a subset of the array to be rendered.

In general, it would be great if axis_item_labels could accept Penzai named arrays/Jax arrays (and internally cast each to a list or equivalent so the object is json-serializable).

import treescope
from penzai import pz

# dimension-dependent axis-item labels:
# should render different labels for each slice of `offset`
axis_item_labels = pz.nx.ones({"x": 10}) * pz.nx.arange("y", 10)  + pz.nx.arange("offset", 3)
treescope.render_array(pz.nx.ones({"offset": 3, "x": 10, "y": 10}), axis_item_labels=axis_item_labels)

# item labels for x and y:
treescope.render_array(pz.nx.ones({"x": 10, "y": 10}), axis_item_labels=[pz.nx.arange("x", 10), pz.nx.arange("y", 10)]) 

[danieldjohnson]

Thanks for the suggestions, I recently ran into a similar situation myself!

I think allowing axis item labels to depend on multiple dimensions would definitely be possible. Unfortunately, I don't think JAX lets you store strings in arrays, so using arrays to specify axis labels seems too limiting.

Another option would be to specify the keys of axis_item_labels as tuples, and allow the values to be nested lists? Perhaps something like

treescope.render_array(..., axis_item_labels={
    "x": ["a", "b", "c"],
    ("y", "z"): [["A", "B", "C"], ["D", "E", "F"]],
})

which would associate each y-z pair with a label. (This would probably require changes to how the axis item labels are actually shown, though.)

[amifalk]

That looks great! Since numpy arrays do support strings, having that as an acceptable input could be a nice-to-have too, but I could always write some internal wrapper if not.