[WIP] using ML optimizers such as Adam in JAX

[lukasheinrich]

Description

this is an exploratory branch to test the ability to use ML optimizers for pyhf (@ncsmith also has some some experiments)

so far this works on a 1Lbb tests case

This crucially also makes the entire fit into a JIT-able function

        loop_state = jax.lax.while_loop(cond,body,loop_state)

which is important for neos (@phi_nate)

$  pyhf cls patched.json --optimizer scipy
/Users/lukasheinrich/Code/pyhfdev/dev/pyhfsrc/src/pyhf/tensor/numpy_backend.py:340: RuntimeWarning: invalid value encountered in log
  return n * np.log(lam) - lam - gammaln(n + 1.0)
{
    "CLs_exp": [
        1.9471461326803052e-09,
        1.301338481151021e-07,
        7.283482301168582e-06,
        0.00029221760192660946,
        0.006624059363217337
    ],
    "CLs_obs": 3.624761348652247e-06
}
$ pyhf cls patched.json --backend jax --optimizer customjax
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
{
    "CLs_exp": [
        1.947135853388778e-09,
        1.3013326218499188e-07,
        7.283455112436134e-06,
        0.00029221673266762734,
        0.006624044552118498
    ],
    "CLs_obs": 3.624759105754125e-06
}

[review-notebook-app]

Check out this pull request on 

See visual diffs & provide feedback on Jupyter Notebooks.

---

Powered by ReviewNB