Explitly trigger jit recompilation

[epignatelli]

After jitting a python function, Is there any way to manually re-jit of the inner, "unjitted", function?

Some use cases are when decorating an intance method with jax.jit, where the first argument is marked as static.
If any of the property of the class is modified, recompilation is not triggered, which I hope is expected.
What if I were to purpously change a property and require recompilation, but avoid nesting jits (i.e., jitting the jitted function)?

E.g.

class Properties(NamedTuple):
    a: int
    b: int


class Foo:
    def __init__(self, properties: Properties):
        self.properties = properties

    @partial(jax.jit, static_argnums=0)
    def update(self, x):
        return self.properties.a / self.properties.b * x


class Bar(Foo):
    def __init__(self, properties):
        super().__init__(properties)
        self.properties = self.properties._replace(a=3)


x = jax.numpy.array((2, 3))
properties = Properties(1, 2)

foo = Foo(properties)
y = foo.update(x)
assert jax.numpy.array_equal(y, 1 / 2 * x)
print(y)

bar = Bar(properties)
y = bar.update(x)
assert jax.numpy.array_equal(y, 3 / 2 * x)  # fails
print(y)

There might be design constraints for which properties cannot be expanded into Foo.

The question is whether there is something like:
self.update = jax.jit(self.update.__inner_python_function__).

[YouJiacheng]

assert jax.numpy.array_equal(y, 3 / 2 * x)  # fails
y = bar.update(x)
print(y)

should be

y = bar.update(x)
assert jax.numpy.array_equal(y, 3 / 2 * x)  # fails
print(y)

And it seems that you doesn't use _replace correctly, since NamedTuple is immutable.

from functools import partial
from typing import NamedTuple

import jax

class Properties(NamedTuple):
    a: int
    b: int


class Foo:
    def __init__(self, properties: Properties):
        self.properties = properties

    @partial(jax.jit, static_argnums=0)
    def update(self, x):
        return self.properties.a / self.properties.b * x


class Bar(Foo):
    def __init__(self, properties):
        super().__init__(properties)
        self.properties = self.properties._replace(a=3)


x = jax.numpy.array((2, 3))
properties = Properties(1, 2)

foo = Foo(properties)
y = foo.update(x)
assert jax.numpy.array_equal(y, 1 / 2 * x)

bar = Bar(properties)
y = bar.update(x)
assert jax.numpy.array_equal(y, 3 / 2 * x) # pass

[patrick-kidger]

In addition to the existing (correct) answer: if you want to use classes around jax.jit then it's usually worth making your class be a PyTree. That way JAX can handle things in the appropriate manner. See Equinox, for example.