Customized `selector` functions for user-defined `Pytree` classes about penzai HOT 3 CLOSED

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.

from penzai.

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.

from penzai.

Closing, as I have several ideas to think about.

from penzai.