Update docs

README.md

@@ -4,15 +4,15 @@
   </a> 
 </p>
 
-<h2 align='center'>Predictive coding networks in JAX.</h2>
+<h2 align='center'>🧠 Predictive coding networks in JAX ⚡️</h2>
 
 JPC is a [JAX](https://github.com/google/jax) library to train neural networks 
 with predictive coding. It is built on top of three main libraries:
 
 * [Equinox](https://github.com/patrick-kidger/equinox), to define neural 
-networks with PyTorch-like syntax, and
+networks with PyTorch-like syntax,
 * [Diffrax](https://github.com/patrick-kidger/diffrax), to solve the PC 
-activity (inference) dynamics.
+activity (inference) dynamics, and
 * [Optax](https://github.com/google-deepmind/optax), for parameter optimisation.
 
 JPC provides a simple but flexible API for research of PCNs compatible with
@@ -28,7 +28,7 @@ useful JAX transforms such as `vmap` and `jit`.
 * [Basic usage](#basic-usage)
 * [Advanced usage](#advanced-usage)
 
-## Installation
+## ️💻 Installation
 
 ```
 pip install jpc
@@ -39,10 +39,10 @@ Requires Python 3.9+, JAX 0.4.23+, [Equinox](https://github.com/patrick-kidger/e
 [Optax](https://github.com/google-deepmind/optax) 0.2.2+, and 
 [Jaxtyping](https://github.com/patrick-kidger/jaxtyping) 0.2.24+.
 
-## Documentation
+## 📖 Documentation
 Available at https://github.com/thebuckleylab.githhub.io/jpc.
 
-## Quick example
+## ⚡️ Quick example
 
 Given a neural network with callable layers
 ```py

jpc/_utils.py

@@ -7,13 +7,13 @@
 from typing import Callable, Optional
 
 
-def get_fc_network(
+def make_mlp(
         key: PRNGKeyArray,
         layer_sizes: PyTree[int],
         act_fn: str,
         use_bias: bool = True
 ) -> PyTree[Callable]:
-    """Defines a fully connected network compatible with predictive coding updates.
+    """Creates a multi-layer perceptron compatible with predictive coding updates.
 
     **Main arguments:**
 

mkdocs.yml

@@ -94,20 +94,20 @@ plugins:
 nav:
     - 'index.md'
     - Predictive coding primer: 'pc_overview.md'
-    - How it works:
+    - ⚙️ How it works:
         - Basic usage: 'basic_usage.md'
         - Advanced usage: 'advanced_usage.md'
         - Extending JPC: 'extending_jpc.md'
-    - Examples:
+    - 📚 Examples:
         - Discriminative PC: 'examples/discriminative_pc.ipynb'
         - Supervised generative PC: 'examples/supervised_generative_pc.ipynb'
         - Unsupervised generative PC: 'examples/unsupervised_generative_pc.ipynb'
         - Hybrid PC: 'examples/hybrid_pc.ipynb'
-    - Basic API:
+    - 🌱 Basic API:
           - 'api/Training.md'
           - 'api/Testing.md'
           - 'api/get_fc_network.md'
-    - Advanced API:
+    - 🚀 Advanced API:
         - 'api/Initialisation.md'
         - 'api/pc_energy_fn.md'
         - 'api/solve_pc_activities.md'