A functionally refactored, differentiable, jax-only pyhf

[phinate]

This is me testing GH discussions :)

So, something I briefly talked about with @lukasheinrich is the idea of a jax-only pyhf implementation. I think this could have some advantages:

• Doesn't force restructuring of current pyhf to make everything differentiable
  • e.g. could natively include differentiable optimizers like in neos without bloating what already exists in pyhf.optimize
  • pull requests such as [WIP] differentiable model construction #742 don't have to be merged in
• Could be more aggressive with jax.vmap and jax.jit without the multiple backend problem
• More things could be jitted if one uses a functional approach, e.g. I managed to jit a dummy version of pyhf.Model.logpdf by implementing something 'pure' by jax standards that had signature logpdf(model, pars, data)(but im pretty sure one could do that with a class method too)

totally just an idea, but I could put some proper time into this if people would be willing to answer all my stupid questions about the current implementation :p

[lukasheinrich]

I think one of the keys questions is:

1. how much of an upside are we missing due to not going "all in" on jax
2. how can we ensure we keep usability from all the backends we have.. upcoming metadiff code should of course help here.

it is true that e.g. batched models which took a huge amount of work in #503 could've been achieved quite easily with jax.vmap

the slowest part in the function evaluation are the interoplators: I put some initial studies here:

https://github.com/lukasheinrich/pyhf-benchmarks/blob/master/colab/jax_benchmarks.ipynb

generally I would expect an upside to appear in places where the JIT can e.g. avoid temporaries (such as those einsum results) by jitting them away

there is also a "gradual" backend idea in which define very specific "levels" in which a backend can appear..

• low-level: just implement tensors and basic stats: pyhf does everything else, build the pyhf model, manaage the logpdf evaluation, implpement modifiers.. this is where we are nw
• mid-level: implement API to consume pyhf spec and have e.g. lhood evaluation, let pyhf do optimization / hypotests
• high-level: implement API to produce full test stat evaluation, best fit results / toys results etc: this could be an API that could even possibly be implemented by ROOT (see notebooks @matthewfeickert cleaned up for fix: Report expected p-values based on their quantiles under background hypotheses #1162 ) this API pyhf doesn't know anything about optimizers, but is given best-fit points, MLE, correlation matrices from a black box etc.. it's the API consumed by calculators / functions-as-a-service. This is also probably the API cabinetry would consume use (@alexander-held )

[lukasheinrich]

checking a single interp code (code1)

re: establishing whether or not we can gain significant upside: it seems like our backend-agnostic one is not too shabby

def func_b(h,a):
    N = h.shape[0]
    M = h.shape[1]
    O = h.shape[3]
    B = a.shape[1]
    
    with jax.experimental.loops.Scope() as s:
        s.out = jnp.zeros((N,M,B,O))
        for i in s.range(N):
            for j in s.range(M):
                for k in s.range(O):
                    delta_up = h[i,j,2,k]/h[i,j,1,k]
                    delta_down = h[i,j,0,k]/h[i,j,1,k]
                    for l in s.range(B):
                        alpha = a[i,l]
                        update = jnp.where(alpha > 0,jnp.power(delta_up, alpha),jnp.power(delta_down, (-alpha)))
                        s.out = jax.ops.index_update(s.out, (i,j,l,k),update)
    return s.out


h = jnp.array(np.random.normal(size = (10,100,3,100)))
a = jnp.array(np.random.normal(size = (10,1)))


func_b = jax.jit(func_b)
start = time.time()
func_b(h,a);
print(time.time()-start)

[lukasheinrich]

but maybe this is something in which jax.vmap will be interesting for some of those loops @phinate

[lukasheinrich]

though for some it seem kinda dependent on the shapes

[lukasheinrich]

maybe a more general comment on "functional refactoring" . I would stay that pyhf currently is not too object-y. The interpolators are essentially just functors while the Model just holds a few data items and provides the logpdf method..

[phinate]

yeah I think so too actually, I was conflating the existence of classes with stateful mutation

[lukasheinrich]

leaving this here:

from numba import jit
import numpy as np
import pyhf
import jax
import jax.experimental.loops
import jax.numpy as jnp

pyhf.set_backend('jax')

@jit(nopython=True)
def numba_func(h,a):
    N = h.shape[0]
    M = h.shape[1]
    O = h.shape[3]
    B = a.shape[1]
    
    out = np.zeros((N,M,B,O))
    for i in np.arange(N):
        for j in np.arange(M):
            for k in np.arange(O):
                delta_up = h[i,j,2,k]/h[i,j,1,k]
                delta_down = h[i,j,0,k]/h[i,j,1,k]
                for l in np.arange(B):
                    alpha = a[i,l]
                    update = delta_up ** alpha if alpha > 0 else delta_down ** (-alpha)
                    out[(i,j,l,k)] = update
    return out


def jax_func(h,a):
    N = h.shape[0]
    M = h.shape[1]
    O = h.shape[3]
    B = a.shape[1]
    
    with jax.experimental.loops.Scope() as s:
        s.out = jnp.zeros((N,M,B,O))
        for i in s.range(N):
            for j in s.range(M):
                for k in s.range(O):
                    delta_up = h[i,j,2,k]/h[i,j,1,k]
                    delta_down = h[i,j,0,k]/h[i,j,1,k]
                    for l in s.range(B):
                        alpha = a[i,l]
                        update = jnp.where(alpha > 0,jnp.power(delta_up, alpha),jnp.power(delta_down, (-alpha)))
                        s.out = jax.ops.index_update(s.out, (i,j,l,k),update)
    return s.out


h = np.random.normal(size = (10,10,3,1000))
a = np.random.normal(size = (10,1))
numba_func(h,a); #jit


jax_h = jnp.array(h)
jax_a = jnp.array(a)


jax_func = jax.jit(jax_func)
jax_func(jax_h,jax_a);  #jit

%%timeit
numba_func(h,a)

13.4 ms ± 335 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
jax_func(h,a)

6.36 ms ± 303 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%%timeit
pyhf.interpolators.code1(h)(a)

3.47 ms ± 143 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)