dimension-specific vmin and vmax in render_array

[eringrant]

A common use case for me is visualizing an N-d array where one axis represents varying semantics (like loss vs. accuracy), but it is still meaningful to visualize the joint array because of shared coordinates. In those cases, I found that the global arguments for vmin and vmax are limiting because the scales of the different semantic dimension are distinct.

It would be great to be able to specify a dimension-specific value for vmin and vmax.
Adapting an example from the docs:

import numpy as np
import treescope

my_3d_array = np.einsum(
    'ijk,i->ijk',
    np.cos(np.arange(5*6*7).reshape((5,6,7)) * 0.1),
    np.array([1.0, 2.0, 3.0, 4.0, 5.0]),
)

As a possible interface, I could think of adapting this global variant:

treescope.render_array(my_3d_array, vmin=-5.0, vmax=5.0)

to a dimension-specific variant like:

treescope.render_array(
    my_3d_array,
    axis_vmin={
        0: (-1.0, -2.0, -3.0, -4.0, -5.0),
    },
    axis_vmax={
        0: (1.0, 2.0, 3.0, 4.0, 5.0),
    },
)

Is this in scope of render_array?