Fix init

README.md

@@ -25,9 +25,9 @@ Requires Python 3.9+, JAX 0.4.23+, [Equinox](https://github.com/patrick-kidger/e
 ## Documentation
 Available at X.
 
-## Quick example
+## Quickstart
 
-Given a neural network with callable layers, for example defined with
+Given a neural network with callable layers defined with
 [Equinox](https://github.com/patrick-kidger/equinox)
 ```py
 import jax

jpc/__init__.py

@@ -1,22 +1,25 @@
 import importlib.metadata
 
-from core._init import (
+from .core import (
     init_activities_with_ffwd as init_activities_with_ffwd,
     init_activities_from_gaussian as init_activities_from_gaussian,
     init_activities_with_amort as init_activities_with_amort
 )
-from core._energies import (
+from .core import (
     pc_energy_fn as pc_energy_fn,
     hpc_energy_fn as hpc_energy_fn
 )
-from core._infer import solve_pc_activities as solve_pc_activities
-from core._grads import (
+from .core import (
     compute_pc_param_grads as compute_pc_param_grads,
     compute_gen_param_grads as compute_gen_param_grads,
     compute_amort_param_grads as compute_amort_param_grads
 )
+from .core import solve_pc_activities as solve_pc_activities
 
-from ._utils import get_fc_network as get_fc_network
+from ._utils import (
+    get_fc_network as get_fc_network,
+    compute_accuracy as compute_accuracy
+)
 from ._train import (
     make_pc_step as make_pc_step,
     make_hpc_step as make_hpc_step

jpc/core/__init__.py

@@ -0,0 +1,15 @@
+from ._init import (
+    init_activities_with_ffwd as init_activities_with_ffwd,
+    init_activities_from_gaussian as init_activities_from_gaussian,
+    init_activities_with_amort as init_activities_with_amort
+)
+from ._energies import (
+    pc_energy_fn as pc_energy_fn,
+    hpc_energy_fn as hpc_energy_fn
+)
+from ._grads import (
+    compute_pc_param_grads as compute_pc_param_grads,
+    compute_gen_param_grads as compute_gen_param_grads,
+    compute_amort_param_grads as compute_amort_param_grads
+)
+from ._infer import solve_pc_activities as solve_pc_activities