yay

setup.cfg

@@ -35,7 +35,8 @@ project_urls =
 [options]
 packages = find:
 install_requires =
-    numpy>=1.13.3
+    pyhf
+    relaxed==0.0.4
 python_requires = >=3.7
 include_package_data = True
 package_dir =

src/neos/examples/nn_pyhf_cls.py

@@ -25,7 +25,7 @@
 pyhf.default_backend = pyhf.tensor.jax_backend(precision="64b")
 
 
-def neos_pyhf_example(rng=PRNGKey(1)):
+def neos_pyhf_example(rng=PRNGKey(1), maxN=10):
 
     # regression net
     init_random_params, predict = stax.serial(
@@ -100,8 +100,6 @@ def update_and_value(i, opt_state, mu):
 
             yield network, metrics, epoch_time
 
-    maxN = 50  # make me bigger for better results!
-
     # Training
     for i, (network, metrics, epoch_time) in enumerate(train_network(maxN)):
         pass  # print(f"epoch {i}:", f'CLs = {metrics["loss"][-1]}, took {epoch_time}s')

tests/test_package.py

@@ -1,5 +1,10 @@
-import neos as m
+import neos
+from neos.examples import neos_pyhf_example
 
 
 def test_version():
-    assert m.__version__
+    assert neos.__version__
+
+
+def test_workflow():
+    assert neos_pyhf_example()