google-deepmind/penzai

Synchronize sliders across array data in Pytree

femtomc opened this issue · 4 comments

https://penzai.readthedocs.io/en/stable/notebooks/treescope_arrayviz.html#slicing-and-scrubbing-with-sliders presents a great way to setup small visuals inside the printing of nested Pytree instances.

One thing I'd like to figure out: is it possible to synchronize the slider state across all arrays?

In general, my computations produce Pytrees whose interpretation will be invalid if a user looks at a visual where some of the sliders are mismatched on slider index.

I'm not sure if this is out of scope, or how this would work, but curious if there are any thoughts about this.

Hmm, right now each of the renderings are produced independently, so it's a bit tricky to implement something like this. Also, the automatic visualization system has been primarily designed as a "better array repr" for debugging and prototyping (in which case you probably already know which slider values are meaningful).

Could you elaborate more on what you're trying to do?

Sure! This is not a "deal breaker" for me, in the least -- was just curious.

To provide more info:

Imagine a Pytree particle collection (a set of samples, roughly) produced by a sequential Monte Carlo algorithm written in a probabilistic programming language: the particles are a nested Pytree with array leaves -- and it contains information like "this particle was created by invoking this program with these arguments, and this was the joint log probability of the random choices under that the measure that program represents, and these are the set of random samples..."

In other words, when I return that collection to the user of the algorithm, one natural thing would be to say -- hey, I want look at this "slice" (single particle -- possibly a slider for this), and peruse various attributes over time (possibly a slider for this). But it's important to have the data "synced up" for some of the sliders (e.g. the "which particle" slider) -- because if you have the slider on e.g. index 51, and another on index 101 -- that's (what I called) an invalid interpretation (looking at attributes of particles which don't match).

Not a big deal at all if out of scope / not easy -- I can figure out an alternative strategy for the visuals I'm considering.

Closing, as I've figured out an alternative strategy.