jax-ml
Pull harvest out into a separate pkg
Hi!
I’ve begun to use the harvest transform for a bunch of little things — for the purposes of this conversation, let’s say one of these things is a type of runtime debugger.
I started with a version of harvest which was a little more restrictive than the one in Oryx. But now I want to apply my debugger to more complicated programs, with control flow primitives.
now, this just about reaches the limit of what I would be comfortable maintaining separately (judging by the amount of touching internals)
Is there any talk of pulling harvest into JAX directly? (E.g. similar to experimental.checkify)?
We have discussed it and I personally would be open to it. I can re-open the conversation with the other Jax devs.
I will say — some of the functionality I’ve developed could be useful as part of a general “jax.experimental.debug” package — paired with checkify.
The pattern is simple and likely easy for you to guess — but it involves recording state out and pairing it with relevant code information (location source lines, modules, etc). I can show you this sometime ~next week.
I’ve found it to be tremendously useful for debugging some of my generative computations.
I don't think we should put harvest in JAX directly; while it's clearly useful I'm not yet satisfied it's a canonical API. (FWIW jax.experimental isn't meant to be a collection of useful JAX extensions, but rather a temporary staging ground for things that will ultimately either be upstreamed into the main jax namespace or else deleted.)
In practical terms, is the upside just not having to pip install another package / have another package requirement?
Fwiw, I've found some of the function transformations (inverse and ildj) here super useful, and the wheel is only 200kb, so I haven't minded that as a requirement for now.
In practical terms, is the upside just not having to pip install another package / have another package requirement?
Yes, that's basically it. I'm only really relying on harvest -- so seems slightly unnecessary.
There's slightly more advanced reasons (re - dependencies) -- e.g. later on, I'd like to support Oryx models as an optional plug-in for my system (so users can write Oryx models, and automatically get something that satisfies our probabilistic programming interface), and I don't want that to complicate my usage of harvest. That might be something I can figure out myself with some clever dependency wizardry.
@mattjj If you (and Sharad) have a sense that there is something which will replace harvest as a canonical API -- I'm perfectly fine to just continue my current ball and dance and then wait to drop in something months down the road.