Raise better error message for pyhf.exceptions.InvalidPdfData for JAX backend

[kratsg]

>>> import pyhf
>>> pyhf.set_backend("jax")
>>> m = pyhf.simplemodels.hepdata_like([10], [15], [5])
>>> pyhf.infer.mle.fit([12.5], m)

crashes like so

with a possible hint?

This error can occur when a JAX Tracer object is passed to a raw numpy function, or a method on a numpy.ndarray object. You might want to check that you are using jnp together with import jax.numpy as jnp rather than using np via import numpy as np. If this error arises on a line that involves array indexing, like x[idx], it may be that the array being indexed x is a raw numpy.ndarray while the indices idx are a JAX Tracer instance; in that case, you can instead write jax.device_put(x)[idx].

[matthewfeickert]

This also fails

import pyhf

if __name__ == "__main__":
    pyhf.set_backend("numpy")
    model = pyhf.simplemodels.hepdata_like([10], [15], [5])
    observations = [12.5]
    print(pyhf.infer.mle.fit(observations, model))

Eval failed for data [12.5] pars: [1.0, 1.0]
Traceback (most recent call last):
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/pdf.py", line 722, in logpdf
    raise exceptions.InvalidPdfData(
pyhf.exceptions.InvalidPdfData: eval failed as data has len 1 but 2 was expected
Traceback (most recent call last):
  File "/tmp/fail.py", line 8, in <module>
    print(pyhf.infer.mle.fit(observations, model))
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/infer/mle.py", line 125, in fit
    return opt.minimize(
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/optimize/mixins.py", line 160, in minimize
    result = self._internal_minimize(**minimizer_kwargs, options=kwargs)
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/optimize/mixins.py", line 38, in _internal_minimize
    result = self._minimize(
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/optimize/opt_scipy.py", line 87, in _minimize
    return minimizer(
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/_minimize.py", line 627, in minimize
    return _minimize_slsqp(fun, x0, args, jac, bounds,
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/slsqp.py", line 375, in _minimize_slsqp
    sf = _prepare_scalar_function(func, x, jac=jac, args=args, epsilon=eps,
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/optimize.py", line 261, in _prepare_scalar_function
    sf = ScalarFunction(fun, x0, args, grad, hess,
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py", line 136, in __init__
    self._update_fun()
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py", line 226, in _update_fun
    self._update_fun_impl()
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py", line 133, in update_fun
    self.f = fun_wrapped(self.x)
  File "/home/feickert/.pyenv/versions/pyhf-CPU/lib/python3.8/site-packages/scipy/optimize/_differentiable_functions.py", line 130, in fun_wrapped
    return fun(x, *args)
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/optimize/opt_numpy.py", line 30, in func
    return objective(constrained_pars, data, pdf)[0]
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/infer/mle.py", line 47, in twice_nll
    return -2 * pdf.logpdf(pars, data)
  File "/home/feickert/Code/GitHub/pyhf/src/pyhf/pdf.py", line 722, in logpdf
    raise exceptions.InvalidPdfData(
pyhf.exceptions.InvalidPdfData: eval failed as data has len 1 but 2 was expected

as I think it is missing axudata

import pyhf

if __name__ == "__main__":
    pyhf.set_backend("numpy")
    model = pyhf.simplemodels.hepdata_like([10], [15], [5])
    observations = [12.5] + model.config.auxdata
    print(pyhf.infer.mle.fit(observations, model))

[0.         0.89587131]

Seems we need a better error message.

[matthewfeickert]

So it seems we get pretty different error messags brought up across backends:

• NumPy nicely raises

pyhf.exceptions.InvalidPdfData: eval failed as data has len 1 but 2 was expected

• JAX does not

This error can occur when a JAX Tracer object is passed to a raw numpy function, or a method on a numpy.ndarray object. You might want to check that you are using `jnp` together with `import jax.numpy as jnp` rather than using `np` via `import numpy as np`. If this error arises on a line that involves array indexing, like `x[idx]`, it may be that the array being indexed `x` is a raw numpy.ndarray while the indices `idx` are a JAX Tracer instance; in that case, you can instead write `jax.device_put(x)[idx]`.

I'm not sure what to do about this at the moment, but it would be good to figure out a way to make sure our own exceptions get raised to the user.