AdvanceSubtensor Error with JAX sampler

[muunetheus]

When using JAX num_pyro NUTS sampler the following error is produced:

MissingInputError: Input 1 of the graph (indices start from 0), used to compute AdvancedSubtensor1(beta_mat ~ Deterministic, <TensorType(int32, vector)>), was not provided and not given a value. Use the Theano flag exception_verbosity='high', for more information on this error.

This error does not occur when using the standard sampler.

Tested with PYMC3 3.10.0 and 3.11.0.

[twiecki]

Thanks! Can you post a minimal example to reproduce the error?

[muunetheus]

Thanks! Can you post a minimal example to reproduce the error?

Working on getting generic/alternative code that will trigger the error. Unfortunately, I can't share the code that generated that error.

[muunetheus]

I managed to get JAX to compile and run by removing Theano shared variable. I read it from this discussion on discourse.

@twiecki Is this expected or an issues to be resolved?

[twiecki]

@muunetheus Yes, shared variables are currently not supported in JAX, something we need to fix.

[brandonwillard]

I think the only situation in which shared variables weren't working was random sampling, but the others should work.

We definitely need a MWE.

[brandonwillard]

If the only remaining issue involves shared variables, then this issue is a duplicate of #4142.