Support theano.shared in jax_funcify

[twiecki]

Currently using a theano.shared leads to a MissingInputError , see pymc-devs/pymc#4142.

[brandonwillard]

@junpenglao, is there anything similar to shared variables in jax? We can always convert these into standard arrays, but then changing the shared value won't have an effect on the JAXed function.

[brandonwillard]

See my comment in the parent issue for an explanation of the real problem.

For the sake of completion. I've added a direct test for shared variables in #76.

We can change the purpose of this issue to center around true shared variable support (i.e. the ability to shared_var.set_value(...) and affect changes in a JAX-jitted function)—if that's even possible—or we can simply add a warning and close this after merging #76 with said warning. @twiecki, what do you think?

[brandonwillard]

Just to be clear, shared variables work as expected under JAX compilation in Theano (i.e. using theano.function), but the functions returned by jax_funcify are oblivious to shared variables and their functionality. As a matter of fact, jax_funcify should never even see a shared variable.

With this in mind, we shouldn't add any sort of warnings; we should simply let jax_funcify err if it is somehow faced with a shared variable. Likewise, the functionality we're considering here would be a nice feature of jax_funcify, but we can't consider it a requirement, because in that sense it's already fulfilled within the framework that defines shared variables (i.e. Theano and its theano.functions).

In other words, adding shared variable support to sample_numpyro_nuts may require a new framework of its own. FYI: if we had samplers written in Theano using Ops with JAX conversions, this functionality would already be supported, since everything would be compiled via theano.function!

[brandonwillard]

I've added comments to pymc-devs/pymc#4142 that should clarify the situation and explain what a viable solution would look like. In the meantime, I'm going to close this issue.