Repository: scikit-hep/pyhf

Issue:
Add precision setting API to set_backend
# Description

At the moment the current way to set the precision for a backend is by passing a backend object into the `set_backend` API

```python
>>> import pyhf
>>> pyhf.set_backend("pytorch")
>>> pyhf.tensorlib.precision
'32b'
>>> pyhf.set_backend(pyhf.tensor.pytorch_backend(precision="64b"))
>>> pyhf.tensorlib.precision
'64b'
```

it would be nicer if there was the ability to also set the precision through `set_backend` through a `kwarg`

```python
>>> import pyhf
>>> pyhf.set_backend("pytorch", precision="64b")
>>> pyhf.tensorlib.precision
'64b'
```

Also, given the large discrepancies that we see with `32b` precision in the PyTorch backend maybe worth having all backends *default* to setting `64b` precision.

```python
>>> import pyhf
>>> pyhf.set_backend("pytorch")
>>> pyhf.tensorlib.precision
'64b'
```


Discussion:

Resolving patch:
diff --git a/src/pyhf/__init__.py b/src/pyhf/__init__.py
--- a/src/pyhf/__init__.py
+++ b/src/pyhf/__init__.py
@@ -1,7 +1,7 @@
 from .tensor import BackendRetriever as tensor
 from .optimize import OptimizerRetriever as optimize
 from .version import __version__
-from .exceptions import InvalidBackend, InvalidOptimizer
+from .exceptions import InvalidBackend, InvalidOptimizer, Unsupported
 from . import events
 
 tensorlib = None
@@ -35,7 +35,7 @@ def get_backend():
 
 
 @events.register('change_backend')
-def set_backend(backend, custom_optimizer=None):
+def set_backend(backend, custom_optimizer=None, precision=None):
     """
     Set the backend and the associated optimizer
 
@@ -44,16 +44,23 @@ def set_backend(backend, custom_optimizer=None):
         >>> pyhf.set_backend("tensorflow")
         >>> pyhf.tensorlib.name
         'tensorflow'
-        >>> pyhf.set_backend(b"pytorch")
+        >>> pyhf.tensorlib.precision
+        '32b'
+        >>> pyhf.set_backend(b"pytorch", precision="64b")
         >>> pyhf.tensorlib.name
         'pytorch'
+        >>> pyhf.tensorlib.precision
+        '64b'
         >>> pyhf.set_backend(pyhf.tensor.numpy_backend())
         >>> pyhf.tensorlib.name
         'numpy'
+        >>> pyhf.tensorlib.precision
+        '64b'
 
     Args:
         backend (`str` or `pyhf.tensor` backend): One of the supported pyhf backends: NumPy, TensorFlow, PyTorch, and JAX
         custom_optimizer (`pyhf.optimize` optimizer): Optional custom optimizer defined by the user
+        precision (`str`): Floating point precision to use in the backend: ``64b`` or ``32b``. Default is backend dependent.
 
     Returns:
         None
@@ -61,27 +68,45 @@ def set_backend(backend, custom_optimizer=None):
     global tensorlib
     global optimizer
 
+    _supported_precisions = ["32b", "64b"]
+    backend_kwargs = {}
+
+    if isinstance(precision, (str, bytes)):
+        if isinstance(precision, bytes):
+            precision = precision.decode("utf-8")
+        precision = precision.lower()
+
     if isinstance(backend, (str, bytes)):
         if isinstance(backend, bytes):
             backend = backend.decode("utf-8")
         backend = backend.lower()
+
+        if precision is not None:
+            backend_kwargs["precision"] = precision
+
         try:
-            backend = getattr(tensor, "{0:s}_backend".format(backend))()
+            backend = getattr(tensor, f"{backend:s}_backend")(**backend_kwargs)
         except TypeError:
             raise InvalidBackend(
-                "The backend provided is not supported: {0:s}. Select from one of the supported backends: numpy, tensorflow, pytorch".format(
-                    backend
-                )
+                f"The backend provided is not supported: {backend:s}. Select from one of the supported backends: numpy, tensorflow, pytorch"
             )
 
-    _name_supported = getattr(tensor, "{0:s}_backend".format(backend.name))
+    _name_supported = getattr(tensor, f"{backend.name:s}_backend")
     if _name_supported:
         if not isinstance(backend, _name_supported):
             raise AttributeError(
-                "'{0:s}' is not a valid name attribute for backend type {1}\n                 Custom backends must have names unique from supported backends".format(
-                    backend.name, type(backend)
-                )
+                f"'{backend.name:s}' is not a valid name attribute for backend type {type(backend)}\n                 Custom backends must have names unique from supported backends"
             )
+        if backend.precision not in _supported_precisions:
+            raise Unsupported(
+                f"The backend precision provided is not supported: {backend.precision:s}. Select from one of the supported precisions: {', '.join([str(v) for v in _supported_precisions])}"
+            )
+    # If "precision" arg passed, it should always win
+    # If no "precision" arg, defer to tensor backend object API if set there
+    if precision is not None:
+        if backend.precision != precision:
+            backend_kwargs["precision"] = precision
+            backend = getattr(tensor, f"{backend.name:s}_backend")(**backend_kwargs)
 
     # need to determine if the tensorlib changed or the optimizer changed for events
     tensorlib_changed = bool(
@@ -102,9 +127,7 @@ def set_backend(backend, custom_optimizer=None):
                     f"The optimizer provided is not supported: {custom_optimizer}. Select from one of the supported optimizers: scipy, minuit"
                 )
         else:
-            _name_supported = getattr(
-                optimize, "{0:s}_optimizer".format(custom_optimizer.name)
-            )
+            _name_supported = getattr(optimize, f"{custom_optimizer.name:s}_optimizer")
             if _name_supported:
                 if not isinstance(custom_optimizer, _name_supported):
                     raise AttributeError(
diff --git a/src/pyhf/cli/infer.py b/src/pyhf/cli/infer.py
--- a/src/pyhf/cli/infer.py
+++ b/src/pyhf/cli/infer.py
@@ -7,7 +7,7 @@
 from ..utils import EqDelimStringParamType
 from ..infer import hypotest
 from ..workspace import Workspace
-from .. import tensor, get_backend, set_backend, optimize
+from .. import get_backend, set_backend, optimize
 
 log = logging.getLogger(__name__)
 
@@ -87,11 +87,11 @@ def cls(
 
     # set the backend if not NumPy
     if backend in ['pytorch', 'torch']:
-        set_backend(tensor.pytorch_backend(precision='64b'))
+        set_backend("pytorch", precision="64b")
     elif backend in ['tensorflow', 'tf']:
-        set_backend(tensor.tensorflow_backend(precision='64b'))
+        set_backend("tensorflow", precision="64b")
     elif backend in ['jax']:
-        set_backend(tensor.jax_backend())
+        set_backend("jax")
     tensorlib, _ = get_backend()
 
     optconf = {k: v for item in optconf for k, v in item.items()}


Test patch:
diff --git a/tests/test_jit.py b/tests/test_jit.py
--- a/tests/test_jit.py
+++ b/tests/test_jit.py
@@ -10,7 +10,7 @@
 @pytest.mark.parametrize('do_grad', [False, True], ids=['no_grad', 'do_grad'])
 @pytest.mark.parametrize('optimizer', ['scipy', 'minuit'])
 def test_jax_jit(caplog, optimizer, do_grad, do_stitch, return_fitted_val):
-    pyhf.set_backend(pyhf.tensor.jax_backend(precision='64b'), optimizer)
+    pyhf.set_backend("jax", optimizer, precision="64b")
     pdf = pyhf.simplemodels.hepdata_like([50.0], [100.0], [10.0])
     data = pyhf.tensorlib.astensor([125.0] + pdf.config.auxdata)
 
@@ -76,7 +76,7 @@ def test_jax_jit(caplog, optimizer, do_grad, do_stitch, return_fitted_val):
 @pytest.mark.parametrize('do_stitch', [False, True], ids=['no_stitch', 'do_stitch'])
 @pytest.mark.parametrize('do_grad', [False, True], ids=['no_grad', 'do_grad'])
 def test_jax_jit_switch_optimizer(caplog, do_grad, do_stitch, return_fitted_val):
-    pyhf.set_backend(pyhf.tensor.jax_backend(precision='64b'), 'scipy')
+    pyhf.set_backend("jax", "scipy", precision="64b")
     pdf = pyhf.simplemodels.hepdata_like([50.0], [100.0], [10.0])
     data = pyhf.tensorlib.astensor([125.0] + pdf.config.auxdata)
 
@@ -132,7 +132,7 @@ def test_jax_jit_switch_optimizer(caplog, do_grad, do_stitch, return_fitted_val)
 )
 @pytest.mark.parametrize('do_grad', [False, True], ids=['no_grad', 'do_grad'])
 def test_jax_jit_enable_stitching(caplog, do_grad, return_fitted_val):
-    pyhf.set_backend(pyhf.tensor.jax_backend(precision='64b'), 'scipy')
+    pyhf.set_backend("jax", "scipy", precision="64b")
     pdf = pyhf.simplemodels.hepdata_like([50.0], [100.0], [10.0])
     data = pyhf.tensorlib.astensor([125.0] + pdf.config.auxdata)
 
@@ -166,7 +166,7 @@ def test_jax_jit_enable_stitching(caplog, do_grad, return_fitted_val):
 )
 @pytest.mark.parametrize('do_stitch', [False, True], ids=['no_stitch', 'do_stitch'])
 def test_jax_jit_enable_autograd(caplog, do_stitch, return_fitted_val):
-    pyhf.set_backend(pyhf.tensor.jax_backend(precision='64b'), 'scipy')
+    pyhf.set_backend("jax", "scipy", precision="64b")
     pdf = pyhf.simplemodels.hepdata_like([50.0], [100.0], [10.0])
     data = pyhf.tensorlib.astensor([125.0] + pdf.config.auxdata)
 
diff --git a/tests/test_public_api.py b/tests/test_public_api.py
--- a/tests/test_public_api.py
+++ b/tests/test_public_api.py
@@ -32,6 +32,14 @@ def test_set_optimizer_by_string(optimizer_name):
     )
 
 
+@pytest.mark.parametrize("precision_level", ["32b", "64b"])
+def test_set_precision_by_string(precision_level):
+    pyhf.set_backend(pyhf.tensorlib.name, precision=precision_level)
+    assert pyhf.tensorlib.precision == precision_level.lower()
+    pyhf.set_backend(pyhf.tensor.numpy_backend(precision=precision_level))
+    assert pyhf.tensorlib.precision == precision_level.lower()
+
+
 @pytest.mark.parametrize("backend_name", [b"numpy", b"tensorflow", b"pytorch"])
 def test_set_backend_by_bytestring(backend_name):
     pyhf.set_backend(backend_name)
@@ -52,6 +60,22 @@ def test_set_optimizer_by_bytestring(optimizer_name):
     )
 
 
+@pytest.mark.parametrize("precision_level", [b"32b", b"64b"])
+def test_set_precision_by_bytestring(precision_level):
+    pyhf.set_backend(pyhf.tensorlib.name, precision=precision_level)
+    assert pyhf.tensorlib.precision == precision_level.decode("utf-8")
+
+
+@pytest.mark.parametrize("precision_level", ["32b", "64b"])
+def test_set_precision_by_string_wins(precision_level):
+    conflicting_precision = "32b" if precision_level == "64b" else "64b"
+    pyhf.set_backend(
+        pyhf.tensor.numpy_backend(precision=conflicting_precision),
+        precision=precision_level,
+    )
+    assert pyhf.tensorlib.precision == precision_level.lower()
+
+
 @pytest.mark.parametrize("backend_name", ["fail", b"fail"])
 def test_supported_backends(backend_name):
     with pytest.raises(pyhf.exceptions.InvalidBackend):
@@ -64,6 +88,12 @@ def test_supported_optimizers(optimizer_name):
         pyhf.set_backend(pyhf.tensorlib, optimizer_name)
 
 
+@pytest.mark.parametrize("precision_level", ["fail", b"fail"])
+def test_supported_precision(precision_level):
+    with pytest.raises(pyhf.exceptions.Unsupported):
+        pyhf.set_backend("numpy", precision=precision_level)
+
+
 def test_custom_backend_name_supported():
     class custom_backend(object):
         def __init__(self, **kwargs):