Open junpenglao opened 3 years ago
I thought static parameters were effectively Constant
s; is that not the correct indicator already?
Regarding shared values, those don't really exist within the JAX compilation context; they're Constant
s at that point. See #73 for an explanation.
Also, if I recall, static_argnums
only applies to the inputs of the JITed function, which implies that such an indicator would only have relevance during the JAX transpilation process and nowhere else.
I thought static parameters were effectively Constants; is that not the correct indicator already?
I guess those are parameters not actually static parameters (thus they are not Constants), but nonetheless due to the limit of XLA you need to treat them as static at runtime. If we can mark them as static, it gives us a bit of additional flexibility to jit those function.
Also, if I recall, static_argnums only applies to the inputs of the JITed function, which implies that such an indicator would only have relevance during the JAX transpilation process and nowhere else.
Yes that's right.
At a high level, we shouldn't add properties to classes unless they're directly relevant to the concepts/objects being modeled by the classes (e.g. a TensorVariable
class should only concern itself with properties relating tensors and/or variables). This particular addition is far too specific to the JAX transpilation process, and, just like the numerous C-related methods attached to our classes, we can incorporate this information differently.
From a lower level, our class implementations need to remain as simple and "static" as possible. Doing so greatly improves the comprehensibility of our code, since it introduces fewer runtime and downstream logic surprises (e.g. avoiding questions like "What's this field, where did it come from, and how did it get set to this?"). Also, we could leverage some individually small—but cumulatively large—performance advantages from this situation (e.g. __slots__
), especially as graphs scale. See issue #72 for related concerns.
That said, the tag
field is better for this situation; however, it's still not clear to me how we would use this information. Can you give a small example of how/when we could use it to avoid the limitations in jnp.reshape
and others?
That said, the tag field is better for this situation; however, it's still not clear to me how we would use this information. Can you give a small example of how/when we could use it to avoid the limitations in jnp.reshape and others?
Adding it to tag
is a nice compromised.
The way we could use it to avoid the limitation in jnp.reshape
is by marking the shape arg tensors being jax_static
, then theano.function(..., mode='jax')
would also work (i.e., does not gives an error during jax.jit
).
I came across something like this in #631. There is a second problem in that scalar symbolic variables become scalar numpy arrays during execution, and these cannot be used as static arguments for Jax functions, because they are not hashable.
Jax jit requires static inputs for some of the function args (for example,
shape
injnp.reshape
,length
injax.lax.scan
). Currently, if these are symbolic input it will breakjax.jit
in https://github.com/pymc-devs/Theano-PyMC/blob/a9275c3dcc998c8cca5719037e493809b23422ff/theano/sandbox/jax_linker.py#L80I propose we add a property to
TensorVariable
in:and
SharedVariable
Then we can detect the additional
static_argnums
in:For user, they will need to mark these variable by hand for now, for example, we can do the following to make the tests pass: