feat: Add jax backend (#377)

* add JAX backend

Author: lukasheinrich

README.md

@@ -56,6 +56,7 @@ Computational Backends:
 - [x] NumPy
 - [x] PyTorch
 - [x] TensorFlow
+- [x] JAX
 
 Available Optimizers
 

docs/installation.rst

@@ -45,6 +45,13 @@ Install latest stable release from `PyPI <https://pypi.org/project/pyhf/>`__...
 
     pip install pyhf[torch]
 
+... with JAX backend
+++++++++++++++++++++
+
+.. code-block:: console
+
+    pip install pyhf[jax]
+
 ... with all backends
 +++++++++++++++++++++
 
@@ -85,6 +92,13 @@ Install latest development version from `GitHub <https://github.com/scikit-hep/p
 
     pip install --ignore-installed -U "git+https://github.com/scikit-hep/pyhf.git#egg=pyhf[torch]"
 
+... with JAX backend
+++++++++++++++++++++++
+
+.. code-block:: console
+
+    pip install --ignore-installed -U "git+https://github.com/scikit-hep/pyhf.git#egg=pyhf[jax]"
+
 ... with all backends
 +++++++++++++++++++++
 

setup.py

@@ -11,13 +11,15 @@
 extras_require = {
     'tensorflow': ['tensorflow~=2.0', 'tensorflow-probability~=0.8'],
     'torch': ['torch~=1.2'],
+    'jax': ['jax~=0.1,>0.1.51', 'jaxlib~=0.1,>0.1.33'],
     'xmlio': ['uproot'],
     'minuit': ['iminuit'],
 }
 extras_require['backends'] = sorted(
     set(
         extras_require['tensorflow']
         + extras_require['torch']
+        + extras_require['jax']
         + extras_require['minuit']
     )
 )

src/pyhf/__init__.py

@@ -90,6 +90,10 @@ def set_backend(backend, custom_optimizer=None):
             if custom_optimizer
             else optimize.pytorch_optimizer(tensorlib=backend)
         )
+    elif backend.name == 'jax':
+        new_optimizer = (
+            custom_optimizer if custom_optimizer else optimize.jax_optimizer()
+        )
     else:
         new_optimizer = (
             custom_optimizer if custom_optimizer else optimize.scipy_optimizer()

src/pyhf/cli/infer.py

@@ -30,7 +30,7 @@ def cli():
 @click.option('--teststat', type=click.Choice(['q', 'qtilde']), default='qtilde')
 @click.option(
     '--backend',
-    type=click.Choice(['numpy', 'pytorch', 'tensorflow', 'np', 'torch', 'tf']),
+    type=click.Choice(['numpy', 'pytorch', 'tensorflow', 'jax', 'np', 'torch', 'tf']),
     help='The tensor backend used for the calculation.',
     default='numpy',
 )
@@ -69,6 +69,8 @@ def cls(
         set_backend(tensor.pytorch_backend(float='float64'))
     elif backend in ['tensorflow', 'tf']:
         set_backend(tensor.tensorflow_backend(float='float64'))
+    elif backend in ['jax']:
+        set_backend(tensor.jax_backend())
     tensorlib, _ = get_backend()
 
     optconf = {k: v for item in optconf for k, v in item.items()}

src/pyhf/optimize/__init__.py

@@ -12,6 +12,18 @@ def __getattr__(self, name):
             # for autocomplete and dir() calls
             self.scipy_optimizer = scipy_optimizer
             return scipy_optimizer
+        elif name == 'jax_optimizer':
+            try:
+                from .opt_jax import jax_optimizer
+
+                assert jax_optimizer
+                self.jax_optimizer = jax_optimizer
+                return jax_optimizer
+            except ImportError as e:
+                raise exceptions.ImportBackendError(
+                    "There was a problem importing jax. The jax optimizer cannot be used.",
+                    e,
+                )
         elif name == 'pytorch_optimizer':
             try:
                 from .opt_pytorch import pytorch_optimizer

src/pyhf/optimize/opt_jax.py

@@ -0,0 +1,70 @@
+"""JAX Optimizer Backend."""
+
+from .. import get_backend, default_backend
+from ..tensor.common import _TensorViewer
+from .autodiff import AutoDiffOptimizerMixin
+import jax
+
+
+def _final_objective(pars, data, fixed_vals, model, objective, fixed_idx, variable_idx):
+    tensorlib, _ = get_backend()
+    tv = _TensorViewer([fixed_idx, variable_idx])
+    pars = tensorlib.astensor(pars)
+    constrained_pars = tv.stitch([fixed_vals, pars])
+    return objective(constrained_pars, data, model)[0]
+
+
+_jitted_objective_and_grad = jax.jit(
+    jax.value_and_grad(_final_objective), static_argnums=(3, 4, 5, 6)
+)
+
+
+class jax_optimizer(AutoDiffOptimizerMixin):
+    """JAX Optimizer Backend."""
+
+    def setup_minimize(
+        self, objective, data, pdf, init_pars, par_bounds, fixed_vals=None
+    ):
+        """
+        Prepare Minimization for AutoDiff-Optimizer.
+
+        Args:
+            objective: objective function
+            data: observed data
+            pdf: model
+            init_pars: initial parameters
+            par_bounds: parameter boundaries
+            fixed_vals: fixed parameter values
+
+        """
+
+        tensorlib, _ = get_backend()
+        all_idx = default_backend.astensor(range(pdf.config.npars), dtype='int')
+        all_init = default_backend.astensor(init_pars)
+
+        fixed_vals = fixed_vals or []
+        fixed_values = [x[1] for x in fixed_vals]
+        fixed_idx = [x[0] for x in fixed_vals]
+
+        variable_idx = [x for x in all_idx if x not in fixed_idx]
+        variable_init = all_init[variable_idx]
+        variable_bounds = [par_bounds[i] for i in variable_idx]
+
+        tv = _TensorViewer([fixed_idx, variable_idx])
+
+        data = tensorlib.astensor(data)
+        fixed_values_tensor = tensorlib.astensor(fixed_values, dtype='float')
+
+        def func(pars):
+            # need to conver to tuple to make args hashable
+            return _jitted_objective_and_grad(
+                pars,
+                data,
+                fixed_values_tensor,
+                pdf,
+                objective,
+                tuple(fixed_idx),
+                tuple(variable_idx),
+            )
+
+        return tv, fixed_values_tensor, func, variable_init, variable_bounds

src/pyhf/tensor/__init__.py

@@ -10,6 +10,13 @@ def __getattr__(self, name):
             # for autocomplete and dir() calls
             self.numpy_backend = numpy_backend
             return numpy_backend
+        elif name == 'jax_backend':
+            from .jax_backend import jax_backend
+
+            assert jax_backend
+            # for autocomplete and dir() calls
+            self.jax_backend = jax_backend
+            return jax_backend
         elif name == 'pytorch_backend':
             try:
                 from .pytorch_backend import pytorch_backend