[WIP] made whole adaptive workflow jax/jit

[mrshirts]

Converted the entire loop in adaptive to jax/jit.
See discussion on issue #340 for discussion of timings.

Intended for inspection, not merging for now.

Failing because jax/jit isn't loaded via conda (not clear it can be right now, conda-forge install on my machine didn't work, had to pip install).

[proteneer]

This will slow things down dramatically, esp. on the GPU. Do you really need 64bit precision?

[mrshirts]

Probably yes - the algorithm doesn't seem to converge well with 32 bit floats. But I can poke around and see if I can isolate the particular problems in convergence.

[mrshirts]

Still need to work on this . . .

[proteneer]

When you call jax.jit, the first invocation will be exceedingly slow it's compiling the JIT kernels. So if you're benchmarking you should call this more than once to amortize out the jit time.

Futhermore, if the arguments have shapes that contain (None,) in one of the shapes, a recompilation is required due to an XLA requirement for each specialization of (None,) into a known shape.

[mrshirts]

Ah, this explains why the partial jitification works better- then it's calling different functions multiple times. Though only 7-8 times in many cases.

[mrshirts]

I will test running mbar initialization several times with different data sets in the same script to see if that helps.

[proteneer]

I also noticed you're declaring every arg to be static, which means that unless you're passing in the exact args you'll be triggering a recompilation. Why did you have to declare them to be static to start with?

[proteneer]

U_kn probably shouldn't be static if you're trying to pass in results from different runs?

[mrshirts]

I had to declare it static because otherwise, gnorm_sci and gnorm_nr end up as traced arrays instead of concrete arrays (since they are functions of traced arrays), and jit refuses to compile the conditional comparing the two, which decides the branch to take, unless they are declared static.

[mrshirts]

@proteneer had a good suggestion for handling the conditional:

import jax
import jax.numpy as jnp
# can't be JIT'd
def foo(a, b, cond_a, cond_b):
    if cond_a < cond_b:
        return a+b
    else:
        return a-b
# foo_jit = jax.jit(foo) # fails
print(foo(0., 1., 2., 3.))
# can be JIT'd
def bar(a, b, cond_a, cond_b):
    return jnp.where(cond_a < cond_b, a+b, a-b)
bar_jit = jax.jit(bar)
print(bar_jit(0., 1., 2., 3.))

[mrshirts]

So, if we accelerate just the inner loop, then we can pull the conditional out of the accelerated code with little loss of timing - all of the conditionals are just comparing floats and assigning values, so there is little use in accelerating them.

[proteneer]

can you write this as jax.lax.while_loop? Just jit-compiling this for loop while be a nightmare for XLA

[mrshirts]

can you write this as jax.lax.while_loop?

I can take a look, though the fact that it's on the adaptive loop that is generally only called once in pymbar means it won't be as useful to do this.

[proteneer]

maybe it makes sense then to just jit the body of this function? (i.e. move everything inside the for loop to a separate function and just jit that)

[mrshirts]

Ah, yes, that makes a lot of sense given the constraints. I'll try that.

[proteneer]

what's the point of 1.0*?

[mrshirts]

N_k is an int, jit (or maybe Jax, I can't recall) complains if it's not converted to float first. I wasn't able to find the best direct conversion to float function that made jit happy. so I just multiplied by 1.0 to do it automatically.

[mrshirts]

Closing in favor of one that just accelerates the inner loop.