Unable to jit logpdf about stheno

[wesselb]

I see! Hmm, this might be challenging. Dispatch currently heavily leverages types, and the type of a PyTree is somewhat troublesome. You're right that jaxtyping offers a PyTree type, but that type seems to only perform instance checking rather than containing the recursive type definition that we would like. I'll have to think about this! It agree that it would be super useful to support PyTrees.

[wesselb]

With the new version of Plum, I believe it should be possible to use the PyTree type from jaxtyping to get the desired behaviour. I'll open issues in the appropriate places.

[wesselb]

Hey @patel-zeel! Nice to hear from you. :)

The reason why this is failing is because there is some logic going on which checks for missing values, as you've noticed, and this logic unfortunately doens't work well with the JIT.

The recommended solution is to use B.jit instead where import lab.jax as B. B.jit is a thin wrapper around jax.jit which runs an additional "compilation step" that takes care of code like the checking for NaNs:

import jax
import jax.numpy as jnp

from stheno.jax import GP, EQ
import lab.jax as B

x = jnp.arange(10)
y = jnp.arange(10)
lengthscale = jnp.array(1.0)
loss_fn = lambda lengthscale: GP(EQ().stretch(lengthscale))(x).logpdf(y)
grad_fn = B.jit(jax.grad(loss_fn))
grad_fn(lengthscale)

Output:

DeviceArray(50.12888031, dtype=float64)

Would you be able to check if this also works on your end?

Thanks for mentioning Chex! I wasn't aware of the library. It looks useful—I'm going to have a closer look!

[patel-zeel]

Thank you, @wesselb, for a quick response! This works, and the depth of your customization is amazingly unimaginable :)

When I try with a slightly different variant (passing a dictionary instead of value), it throws plum.function.NotFoundLookupError to me:

import jax
import jax.numpy as jnp

from stheno.jax import GP, EQ
import lab.jax as B

x = jnp.arange(10)
y = jnp.arange(10)
params = {"lengthscale": jnp.array(1.0)}
loss_fn = lambda params: GP(EQ().stretch(params["lengthscale"]))(x).logpdf(y)
grad_fn = B.jit(jax.grad(loss_fn))
print(grad_fn(params))

Output

Traceback (most recent call last):
  File "/home/patel_zeel/gpax/testbed.py", line 12, in <module>
    print(grad_fn(params))
  File "/home/patel_zeel/miniconda3/envs/ajax/lib/python3.9/site-packages/lab/generic.py", line 146, in __call__
    return _jit_run(
  File "/home/patel_zeel/miniconda3/envs/ajax/lib/python3.9/site-packages/plum/function.py", line 591, in __call__
    method, return_type = self.resolve_method(*sig_types)
  File "/home/patel_zeel/miniconda3/envs/ajax/lib/python3.9/site-packages/plum/function.py", line 556, in resolve_method
    method, return_type = self._methods[self.resolve_signature(signature)]
  File "/home/patel_zeel/miniconda3/envs/ajax/lib/python3.9/site-packages/plum/function.py", line 492, in resolve_signature
    raise NotFoundLookupError(
plum.function.NotFoundLookupError: For function "_jit_run", signature Signature(builtins.function, builtins.dict, builtins.dict, builtins.dict) could not be resolved.

Also, I wonder how to use B.jit at places like jax.lax.scan function, which applies jax.jit internally.
Update: I checked with jax.lax.scan, and it works! It looks like lax.scan does not explicitly use jax.jit internally. Now only plum.function.NotFoundLookupError problem remains.

[wesselb]

I think the problem now is that is typing is a little restrictive. In particular, a dictionary with JAX-valued values isn't recognised as a JAX object, and that's where the method error comes from. It is possible to add that method manually, as follows:

from types import FunctionType

import jax
import jax.numpy as jnp

import lab.jax as B
from plum import Union, Dict
from stheno.jax import GP, EQ


@B.generic._jit_run.dispatch
def _jit_run(
    f: FunctionType,
    compilation_cache: dict,
    jit_kw_args: dict,
    *args: Dict[object, B.JAXNumeric],
    **kw_args,
):
    return B.generic._jit_run.invoke(
        FunctionType, dict, dict, B.JAXNumeric,
    )(f, compilation_cache, jit_kw_args, *args, **kw_args)


x = jnp.arange(10)
y = jnp.arange(10)
params = {"lengthscale": jnp.array(1.0)}
loss_fn = lambda params: GP(EQ().stretch(params["lengthscale"]))(x).logpdf(y)
grad_fn = B.jit(jax.grad(loss_fn))
print(grad_fn(params))

{'lengthscale': DeviceArray(50.12888031, dtype=float64)}

It would be possible to make a small amendment to LAB so that this is the default behaviour, if that would be desirable. :)

[patel-zeel]

Thanks for the solution, @wesselb. Due to the PyTree concept gaining popularity in JAX, dictionaries of parameters or lists and tuples of parameters or a mixture of these are also used sometimes as parameters. Is it possible to make B.jit work with any PyTree? Maybe PyTree from jaxtyping could be used somehow, but I am not sure how.

[patel-zeel]

Yes, from the links you have shared, it looks like a harder problem. Maybe there can be a hotfix specifically for lab and Stheno till a more generic solution is available.

[wesselb]

A hacky halfway house solution is to check for JAX-like objects in the first 10 (or so) layers of a PyTree... Then it wouldn't detect things like ((((((((((jnp.array(1),)))))))))), but maybe that's okay as a temporary solution. Do you think that would be reasonable? Or do you perhaps have another fix in mind?

[patel-zeel]

I think 10 or so layers seem practical for most applications. However, how difficult would it be to convert that static number 10 to a dynamic depth by recursively checking PyTree?

[wesselb]

I think that it would be possible to convert the static number to a dynamic depth. However, perhaps the right solution here is to see if we can actually give PyTrees first-class support. I’ll soon be working on a 2.0 of Plum, which is where currently the restrictions derive from. I will put PyTree support on the list of desired improvements!