Passing functions or classes to scan or best equivalent way to achieve the same thing

[mchagneux]

Hi all,

I'm porting to Jax one of my Pytorch codebase where one big function (which I'll denote F) takes:

• some parameters (in a dict) which I'll denote Theta.
• some input data (which is actually a sequence)

and returns some quantity of interest. The inner workings of that function involve a for-loop on the data with some carry-over.

Up to now I managed to get everything working with lax.scan and putting Theta in the carry, when Theta is a Pytree with np.array leaves. However one convenient aspect of my Pytorch implementation is that some elements of Theta are actually nn modules (which contain their own parameters but can also be called at some point in F). I can give any module I want (which is practical for me because these modules correspond to arbitrarily complex parameterized functions on which I want to run experiments).

I wanted to reproduce this behaviour by passing some tuples of (parameters, function _to_apply) in the carry of lax.scan, or even passing some Haiku model, but apparently this is not allowed (I get error messages saying that "function" is not a valid type for the carry of scan).

I can get around this by passing some "string" in the carry with the name of the function to apply and keeping a dict of these functions somewhere and using e.g lax.cond, but that seems a big ugly so I was wondering if there was some other way ?

Thanks in advance!

[soraros]

I think something like this will work:

from jax import jit
from jax.tree_util import Partial

th = (Partial(lambda params, x: params['x'] + x), dict(x=1))

@jit
def app(t, x):
  f, params = t
  return f(params, x)

app(th, 1)

[YouJiacheng]

FYI, if you want to hide everything, I think you can directly use:

from jax.tree_util import Partial
theta = Partial(lambda params, x: params['x'] + x, dict(x=1))

instead of using a tuple.
In this way, you still can evaluate grad w.r.t. theta.

def call_func(fun):
    return fun()

def f(k1, k2, b):
    return k1 * k2 + b

from jax.tree_util import Partial
theta = Partial(f, 1.0, 2.0, b=3.0)
fun_grad = jax.grad(call_func)(theta)
print(fun_grad.args)
print(fun_grad.keywords)

output will be:

(DeviceArray(2., dtype=float32, weak_type=True), DeviceArray(1., dtype=float32, weak_type=True))
{'b': DeviceArray(1., dtype=float32, weak_type=True)}

And you can even directly using optax to optimize theta! Just like what you do in torch!

# above code is same
import optax
optimizer = optax.sgd(1e-2)
opt_state = optimizer.init(theta)
update, opt_state = optimizer.update(fun_grad, opt_state)
theta = optax.apply_updates(theta, update)
print(theta.args)
print(theta.keywords)

output will be:

(DeviceArray(0.98, dtype=float32), DeviceArray(1.99, dtype=float32))
{'b': DeviceArray(2.99, dtype=float32)}

Amazing!
Note function itself (theta.func) will not be evaluated grad or updated, even if it is a closure or functools.partial, since it will be treated as aux_data in tree_flatten.

[shoyer]

If you like this sort of thing, you might also be interested in Equinox for writing neural net layers: https://github.com/patrick-kidger/equinox

[mchagneux]

Sounds really practical ! I was actually planning to instantiate Haiku modules then pass the (renamed) hk.transform-ed functions into tree_util.Partial, then go ahead with my code as unaltered as possible, but this makes it conceptually simpler, I'll have a try!

Any notes on performance ? E.g. am I guaranteed that a forward pass through equinox.nn.Linear will be as fast as a forward pass through, say, haiku.nn.Linear ? I prefer to ask because sometimes more largely used libraries benefit from clever implementation tricks under the hood, which in the end matter a lot.

[shoyer]

There are no special performance tricks. All of these libraries should have identical performance under JIT.

[mchagneux]

Ok thanks.