Is it possible to pass Type[class] as parameter for a jit'ed function?

[quattro]

Hi all,

Thanks again for providing such a powerful framework for computing/optimization. I have some complex code that takes as input a class type, to enable flexibility in performing some inference. Specifically, I've ported over some of the Trust Region optimization from scipy for my project and have a function that looks like,

def infer_trustregion(model: BaseModel,
                      cls: Type[BaseQuadraticProblem] = CGSteihaugSubproblem,
                      options: TrustRegionOptions = None) -> AlphaResults:

    some_params =model.init()
    sub_p = cls(some_params)
    [...]

I've already extended my BaseModel class to implement PyTree interfaces to allow jit functionality, but what is less clear to me is how enable jit of Type[BaseQuadraticProblem]. The reason I'd like to jit this level of functionality, rather than after instantiating using the cls parameter, is that some a single gradient-like step is performed by model to initialize the BaseQuadraticProblem.

One work-around would be to enable jit on for init separately and then push jit to sub-calls within this function, but this seems sub-optimal rather than jiting the entire function.

Is it possible in jax to jit a function with a Type[...] parameter at the moment? I apologize if this is in the documentation somewhere and I missed it.

[jakevdp]

Yes, this is possible as long as the type argument is declared as static.

Additionally, if you want a jitted function to return a custom object, you need to tell jax how to convert it to a pytree (basically a structure of arrays representing the object).

Here's a quick example of creating a custom type, registering it with JAX, and then passing the type to a jitted function and returning the constructed object:

from jax import jit, partial
from jax.tree_util import register_pytree_node_class

@register_pytree_node_class
class Model:
  def tree_flatten(self):
    # specify how to serialize model into a JAX pytree
    return ((self.x,), None)

  @classmethod
  def tree_unflatten(cls, aux_data, children):
    # specify how to build a model from a JAX pytree
    return cls(*children)

  def __init__(self, x):
    self.x = x

  def __repr__(self):
    return f"Model({self.x})"

# Model class must be static.
@partial(jit, static_argnums=0)
def make_model(cls, y):
  x = y ** 2
  return cls(x)

print(make_model(Model, 2))
# Model(4)

Hopefully you can modify this approach to your own application!

[quattro]

Fantastic! Thanks for the speedy reply. My current implementation of BaseModel is indeed static. It's essentially a glorified NamedTuple with an interface to compute a likelihood and retrieve gradients.