Customized `selector` functions for user-defined `Pytree` classes

[femtomc]

Hi! Awesome package -- had a great sense of "I absolutely must use this" when I was looking at the documentation and the repository.

I have a library I'm developing for probabilistic programming, where I define custom nested Pytree classes. The "getter" and "setter" interfaces for these classes are not just like normal Pytree accessing (via fields) -- they are defined to ensure consistency with the rest of the semantics of the system.

I'd like to use penzai and penzai.treescope with these objects -- so I'm wondering if there's some way to overload the selectors to use my own interfaces, and have this reflected in e.g. the lambda functions to access pieces of data from the visuals?

Thanks -- and again, awesome package -- I'm totally willing to read your code, and try to figure out a solution to the above, if you're enthusiastic about it.

[danieldjohnson]

Thanks for the question!

I'm not sure I fully understand what you are trying to do? By design, Penzai's selector system uses the actual jax.tree_util flattening and unflattening operations, so if you have custom tree flattening and unflattening operations, the selectors should compose with them correctly. You can't overload selectors with the current design, because then it would potentially go out of sync with how JAX sees the data structures under the hood.

penzai.treescope, on the other hand, is primarily a pretty-printer, and is somewhat separate from the selector system. You can implement your own handler for custom objects, including changing what gets copied when you click the "copy path" button. Doing so is a bit complex right now, but you might be able to figure it out by looking at some of the handlers in the handlers directory of penzai/treescope. I'd also like to make extending treescope easier in the future, and having specific use cases might help me figure out the best way to do that!

That said, I'm curious what problems you're running into right now? Have you run into problems trying to use Penzai with your custom classes already? There's a chance things would "just work" out of the box.

[femtomc]

Thanks for the answer!

Right, a bit tough to explain succinctly, but I'll give it a shot: the data structures which I'm letting users pretty print have a specified interface for accessors and setters. This is different than just the Pytree flattening / unflattening operators and what you can build on top of them -- I hide a lot of the complexity of changing the data, as well as gaining access to the "semantically relevant" pieces of data in these objects.

These objects can represent things like "these are the samples that a probabilistic program sampled when it was executed" -- they can also represent things like "the constraints that a probabilistic program should execute under". In both cases, how these objects are constructed is important, and I have something very similar to a "data DSL" for creating these objects, and mutating them.

Maybe an immediate question (and my attempt at a quick answer): why have a "data DSL" at all -- and not just have a system like the selector system which works on Pytree? Because in the programs written in my system, "the identity of a random variable and its sampled value" is specified using a type of addressing system which is different from where the data actually lives in the fields of the Pytree. That's the main problem -- and the thing I'm trying to resolve (by e.g. providing a custom handler so that -- when a user pretty prints with penzai, the "copy" buttons use our addressing system, which also shows them where that random variable lives in the code that they wrote).

From your answer, I think I should investigate:

You can implement your own handler for custom objects, including changing what gets copied when you click the "copy path" button.

Just to clarify -- penzai is working perfectly fine (according to its design principles!) with my objects! I'm just hoping to provide a bit more of a walled garden for some of my objects.

[femtomc]

Closing, as I have several ideas to think about.