### Install Development Dependencies

Source: https://github.com/jaxleyverse/jaxley/blob/main/CONTRIBUTING.md

Install the project using pip with development and documentation dependencies. This command should be run after cloning the repository.

```bash
pip install -e ".[dev, doc]"
```

--------------------------------

### Install Build and Twine

Source: https://github.com/jaxleyverse/jaxley/wiki/Release-workflow

Install the necessary tools for building and uploading Python packages to PyPI. This is part of the fallback release procedure.

```shell
pip install build twine
```

--------------------------------

### Install Jaxley with CPU Support

Source: https://github.com/jaxleyverse/jaxley/blob/main/README.md

Installs the Jaxley library with default CPU support using pip.

```sh
pip install jaxley
```

--------------------------------

### Import Optax Optimizer

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Import the optax library for optimization algorithms. Ensure optax is installed via 'pip install optax'.

```python
import optax
```

--------------------------------

### Build and Serve Sphinx Documentation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/README.md

Run this command to build the Sphinx documentation from Jupyter notebooks and serve it locally. Ensure you have Sphinx and its dependencies installed.

```bash
make html && python -m http.server --directory _build/html
```

--------------------------------

### Initial Parameter Sharing Configuration

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/10_advanced_parameter_sharing.ipynb

Demonstrates an initial setup for parameter sharing where the same parameter is used for all synapses, or individual parameters are made trainable. This serves as a baseline for more advanced sharing techniques.

```python
net = ...  # See tutorial on Basics of Jaxley.

# The same parameter for all synapses
net.make_trainable("Ionotropic_gS")

# An individual parameter for every synapse.
net.select(edges="all").make_trainable("Ionotropic_gS")

# Share synaptic conductances emerging from the same neurons.
net.copy_node_property_to_edges("cell_index")
sub_net = net.select(edges=[0, 1, 2])
sub_net.edges["controlled_by_param"] = sub_net.edges["pre_global_cell_index"]
sub_net.make_trainable("Ionotropic_gS")
```

--------------------------------

### Parameter Transformation Setup

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Imports the necessary Jaxley optimize transforms module. This is required for defining custom parameter transformations.

```python
import jaxley.optimize.transforms as jt
```

--------------------------------

### Install Jaxley with GPU Support

Source: https://github.com/jaxleyverse/jaxley/blob/main/README.md

Installs Jaxley with GPU support by first installing JAX with CUDA support. This example uses CUDA 13.

```sh
pip install -U "jax[cuda13]"
```

--------------------------------

### Install Jaxley from Source

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/installation.md

Installs Jaxley in editable mode from its source repository. Requires pip version 21.3 or higher.

```sh
git clone https://github.com/jaxleyverse/jaxley.git
cd jaxley
pip install -e .
```

--------------------------------

### Jaxley Network Setup and Imports

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/10_advanced_parameter_sharing.ipynb

Imports necessary Jaxley modules and defines network components for building a network. This is boilerplate for network construction.

```python
import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.connect import fully_connect
from jaxley.synapses import IonotropicSynapse
```

--------------------------------

### Run Integration with Precomputed Transition Matrix

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/how_to_guide/choose_solver.ipynb

After precomputing the transition matrix for the 'exp_euler' solver, subsequent calls to `jx.integrate` for the same cell configuration will automatically utilize the precomputed matrix, leading to faster simulations. This example demonstrates running integration twice after setting the transition matrix.

```python
v1 = jx.integrate(cell, solver="exp_euler")
cell.set("gLeak", 1e-4)
v2 = jx.integrate(cell, solver="exp_euler")
```

--------------------------------

### Integrate Network Simulation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/05_channel_and_synapse_models.ipynb

Executes the network simulation using the defined setup, stimulus, and recordings. The result `s` contains the recorded voltage traces.

```python
s = jx.integrate(net)
```

--------------------------------

### Define and Simulate Cell with Ion Dynamics in Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/11_ion_dynamics.ipynb

This snippet demonstrates setting up a cell with specific ion pumps and diffusion mechanisms, then simulating its behavior. Ensure Jaxley and related modules are installed.

```python
import jaxley as jx
from jaxley.pumps import CaPump, CaNernstReversal
from jaxley_mech.channels.l5pc import CaHVA


branch = jx.Branch()
cell = jx.Cell(branch, parents=[-1, 0, 0])

# Insert a voltage-gated calcium channel.
cell.insert(CaHVA())

# Insert a mechanism which modifies the intracellular calcium based on the calcium current.
cell.insert(CaPump())

# Insert a mechanism that updates the calcium reversal potential based on the intracellular calcium level.
cell.insert(CaNernstReversal())

# Let the intracellular calcium diffuse within the cell.
cell.diffuse("CaCon_i")
cell.set("axial_diffusion_CaCon_i", 2.0)

# Record the intracellular calcium concentration and simulate.
cell.record("CaCon_i")
cacon_i = jx.integrate(cell, t_max=100.0, delta_t=0.025)
```

--------------------------------

### Create Jaxley Network Components

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/dev/indexing.ipynb

Initializes a compartment, a branch, a cell, and a network with multiple cells. This sets up the basic structure for indexing examples.

```python
import jaxley as jx

comp = jx.Compartment()
branch = jx.Branch(comp, 1)
cell = jx.Cell(branch, [-1, 0, 0])
net = jx.Network([cell for _ in range(2)])
```

--------------------------------

### Visualize SWC graph with networkX and Matplotlib

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/dev/graph_backend.ipynb

Visualize the imported SWC graph using networkX and Matplotlib. Ensure networkX and Matplotlib are installed.

```python
import networkx as nx
import matplotlib.pyplot as plt

pos = {k: (v["x"], v["y"]) for k, v in swc_graph.nodes.items()}
fig, ax = plt.subplots(1, 1, figsize=(10, 4))
nx.draw(swc_graph, pos=pos, with_labels=True, font_size=7, node_size=150, ax=ax)
```

--------------------------------

### Import Jaxley and Core Libraries

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/11_ion_dynamics.ipynb

Imports necessary libraries for Jaxley simulations, including plotting and numerical computation tools. Ensure these libraries are installed.

```python
import matplotlib.pyplot as plt
import numpy as np
import jax
import jax.numpy as jnp
from jax import jit

import jaxley as jx
from jaxley.channels import Na, K, Leak
```

--------------------------------

### Perform Parameter Sweep with For-Loop

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Run simulations for multiple parameter sets using a standard Python for-loop. This is the least efficient method but demonstrates the basic sweep setup.

```python
# Define 5 sets of sodium and potassium conductances.
all_params = jnp.asarray(np.random.rand(5, 2))

voltages = jnp.asarray([simulate(params) for params in all_params])
print("voltages.shape", voltages.shape)
```

--------------------------------

### Implement Ion Diffusion in Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/how_to_guide/ion_diffusion.ipynb

This snippet demonstrates how to set up a cell with a voltage-gated calcium channel, an ion pump, and a mechanism to update the calcium reversal potential. It then enables intracellular calcium diffusion. Requires `pip install jaxley-mech`.

```python
from jaxley.pumps import CaPump, CaNernstReversal
from jaxley_mech.channels.l5pc import CaHVA


branch = jx.Branch()
cell = jx.Cell(branch, parents=[-1, 0, 0])

# Insert a voltage-gated calcium channel.
cell.insert(CaHVA())

# Insert an ion pump which modifies the intracellular calcium based on the calcium current.
cell.insert(CaPump())

# Insert a mechanism that updates the calcium reversal potential based on the intracellular calcium level.
cell.insert(CaNernstReversal())

# Let the intracellular calcium diffuse within the cell.
cell.diffuse("CaCon_i")
cell.set("axial_resistivity_CaCon_i", 1_000.0)
```

--------------------------------

### Setting up Cell Parameters and Running Simulation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

This code snippet initializes cell parameters, including conductances and leak properties, and then runs a simulation to obtain the observed voltage trace.

```python
cell.compute_compartment_centers()
direct_dists = distance_direct(cell.soma.branch(0).comp(0), cell)
cell.nodes["dist_from_soma"] = direct_dists
gH_conductance = (-0.8696 + 2.087 * np.exp(cell.basal.nodes["dist_from_soma"] * 0.0031)) * 8e-5
cell.basal.set("H_gH", gH_conductance)

cell.set("Leak_gLeak", 3e-05)
cell.set("Leak_eLeak", -75.0)

cell.set("eNa", 50.0)
cell.set("eK", -85.0)
```

```python
x_o = jx.integrate(cell)[0]  # [0] gets rid of the batch-dimension.
```

--------------------------------

### Configure and Run Simulation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/11_ion_dynamics.ipynb

Sets up simulation parameters, clears existing recordings and stimuli, records voltage and intracellular calcium, applies a step current stimulus, and integrates the cell model over time.

```python
t_max = 20.0
dt = 0.025

cell.delete_recordings()
cell.delete_stimuli()

cell.branch(0).comp(0).record("v")
cell.branch(0).comp(0).record("CaCon_i")
cell.branch(0).comp(0).stimulate(jx.step_current(1.0, 10.0, 0.03, dt, t_max))
v_and_ca = jx.integrate(cell, delta_t=dt, t_max=t_max)
```

--------------------------------

### Simulate and Plot Initial Parameters

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

Simulates the model with the initial parameters and plots the voltage over time. This helps visualize how well the initial guess matches the observed data.

```python
v = simulate(initial_params)

fig, ax = plt.subplots(1, 1, figsize=(5, 2))
_ = ax.plot(time_vec, v.T, c="#41ae76")
_ = ax.set_ylim([-90, 60])
_ = ax.set_xlabel("Time (ms)")
_ = ax.set_ylabel("Voltage (mV)")
plt.show()
```

--------------------------------

### Configure Jax for Computation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/01_morph_neurons.ipynb

Set up Jax for 64-bit precision and specify the CPU as the platform for computation. This is essential for numerical stability and performance.

```python
from jax import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")
```

--------------------------------

### Jaxley Configuration and Imports

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/02_small_network.ipynb

Sets up Jax for 64-bit precision and CPU usage, and imports necessary libraries for neural network simulations.

```python
from jax import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import matplotlib.pyplot as plt
import numpy as np
import jax.numpy as jnp
from jax import jit

import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import fully_connect, connect
```

--------------------------------

### Get Parameters for a Group

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/06_groups.ipynb

Retrieves the current trainable parameters, showing a single shared parameter for the 'Na_gNa' conductance across all fast-spiking neurons.

```python
network.get_parameters()
```

--------------------------------

### Save and Load Jaxley Network with Pickle

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/how_to_guide/save_and_load.ipynb

Demonstrates how to save a Jaxley network to a file using pickle and then load it back. Ensure the file path is valid for writing and reading.

```python
import jaxley as jx
import pickle

# ... define network, cell, etc.
network = jx.Network([cell1, cell2])

# Save.
with open("path/to/file.pkl", "wb") as handle:
    pickle.dump(network, handle)

# Load.
with open("path/to/file.pkl", "rb") as handle:
    network = pickle.load(handle)
```

--------------------------------

### Build, Stimulate, Record, and Simulate a Neuron

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/01_morph_neurons.ipynb

This snippet demonstrates the complete workflow of creating a neuron, adding channels, setting parameters, stimulating it, recording its activity, and running a simulation.

```python
import jaxley as jx
from jaxley.channels import Na, K, Leak
import matplotlib.pyplot as plt


# Build the cell.
comp = jx.Compartment()
branch = jx.Branch(comp, ncomp=2)
cell = jx.Cell(branch, parents=[-1, 0, 0, 1, 1])

# Insert channels.
cell.insert(Leak())
cell.branch(0).insert(Na())
cell.branch(0).insert(K())

# Change parameters.
cell.set("axial_resistivity", 200.0)

# Visualize the morphology.
cell.compute_xyz()
fig, ax = plt.subplots(1, 1, figsize=(4, 4))
cell.vis(ax=ax)

# Stimulate.
current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0)
cell.branch(0).loc(0.0).stimulate(current)

# Record.
cell.branch(0).loc(0.0).record("v")

# Simulate and plot.
v = jx.integrate(cell, delta_t=0.025)
plt.plot(v.T)
```

--------------------------------

### Build MkDocs Documentation Locally

Source: https://github.com/jaxleyverse/jaxley/blob/main/mkdocs/README.md

Run this command from the mkdocs subfolder to build the documentation locally. It first converts Jupyter notebooks to markdown.

```bash
jupyter nbconvert --to markdown ../docs/tutorials/*.ipynb --output-dir docs/tutorial/
mkdocs serve
```

--------------------------------

### Modify Group Properties

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/06_groups.ipynb

After defining a group, you can access it as an attribute of the network object. This example shows how to set a property ('radius') for all components within the 'apical' group.

```python
network.apical.set("radius", 0.3)
```

--------------------------------

### Set up Stimulation and Recording

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/05_channel_and_synapse_models.ipynb

Configures the cell for simulation by deleting existing stimuli and recordings, applying a new current stimulus, and setting up voltage recording.

```python
cell.delete_stimuli()
current = jx.step_current(1.0, 1.0, 0.1, 0.025, 10.0)
cell.branch(0).comp(0).stimulate(current)

cell.delete_recordings()
cell.branch(0).comp(0).record()
```

--------------------------------

### Define Custom Synapse Model in Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/05_channel_and_synapse_models.ipynb

Implement a custom synapse by inheriting from `jaxley.synapses.synapse.Synapse` and defining `update_states` and `compute_current` methods. This example defines a simple cholinergic synapse.

```python
import jax.numpy as jnp
from jaxley.synapses.synapse import Synapse


class TestSynapse(Synapse):
    """
    Compute syanptic current and update syanpse state.
    """
    def __init__(self, name = None):
        super().__init__(name)
        self.synapse_params = {"gChol": 0.001, "eChol": 0.0}
        self.synapse_states = {"s_chol": 0.1}

    def update_states(self, states, delta_t, pre_voltage, post_voltage, params):
        """Return updated synapse state and current."""
        s_inf = 1.0 / (1.0 + jnp.exp((-35.0 - pre_voltage) / 10.0))
        exp_term = jnp.exp(-delta_t)
        new_s = states["s_chol"] * exp_term + s_inf * (1.0 - exp_term)
        return {"s_chol": new_s}

    def compute_current(self, states, pre_voltage, post_voltage, params):
        g_syn = params["gChol"] * states["s_chol"]
        return g_syn * (post_voltage - params["eChol"])
```

--------------------------------

### Configure Checkpointing

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Sets up checkpointing parameters for memory-efficient training by defining time points and levels for gradient checkpointing.

```python
t_max = 5.0
dt = 0.025

levels = 2
time_points = t_max // dt + 2
checkpoints = [int(np.ceil(time_points**(1/levels))) for _ in range(levels)]
```

--------------------------------

### Set Ion Name for Pump

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/11_ion_dynamics.ipynb

Specify which ion concentration the pump modifies by setting the `ion_name` attribute. This example shows setting it to 'CaCon_i' for intracellular calcium concentration.

```python
self.ion_name = "CaCon_i"
```

--------------------------------

### Set Up Simulation Parameters

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/02_small_network.ipynb

Define parameters for network stimulation, such as delay, amplitude, and duration. These values are used to configure the input currents for the simulation.

```python
# Stimulus.
i_delay = 3.0  # ms
i_amp = 0.05  # nA
i_dur = 2.0  # ms
```

--------------------------------

### Iterating and Modifying Network Components in Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/00_jaxley_api.ipynb

Demonstrates how to iterate over cells, branches, and compartments in a Jaxley network to modify their properties. This example shows setting radius and length based on conditions.

```python
# We set the radiuses to random values...
radiuses = np.random.rand((24))
net.set("radius", radiuses)

# ...and then we set the length to 100.0 um if the radius is >0.5.
for cell in net:
    for branch in cell:
        for comp in branch:
            if comp.nodes.iloc[0]["radius"] > 0.5:
                comp.set("length", 100.0)

# Show the first five compartments:
net.nodes[["radius", "length"]][:5]
```

--------------------------------

### Assemble Network and Interact with Modules in Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/00_jaxley_api.ipynb

Demonstrates assembling a network from compartments, branches, and cells, then navigating and modifying modules using views and groups. It also shows how to insert channels and connect cells.

```python
import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import connect
import matplotlib.pyplot as plt
import numpy as np


# Assembling different Modules into a Network
comp = jx.Compartment()
branch = jx.Branch(comp, ncomp=1)
cell = jx.Cell(branch, parents=[-1, 0, 0])
net = jx.Network([cell]*3)

# Navigating and inspecting the Modules using Views
cell0 = net.cell(0)
cell0.nodes

# How to group together parts of Modules
net.cell(1).add_to_group("cell1")

# inserting channels in the membrane
with net.cell(0) as cell0:
    cell0.insert(Na())
    cell0.insert(K())

# connecting two cells using a Synapse
pre_comp = cell0.branch(1).comp(0)
post_comp = net.cell1.branch(0).comp(0)

connect(pre_comp, post_comp)
```

--------------------------------

### Configure and Apply Stimuli

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/02_small_network.ipynb

Sets up a step current stimulus and applies it to the first 10 neurons in the input layer. It also clears any existing stimuli and recordings before applying new ones.

```python
current = jx.step_current(i_delay, i_dur, i_amp, dt, t_max)
net.delete_stimuli()
for stim_ind in range(10):
    net.cell(stim_ind).branch(0).loc(0.0).stimulate(current)

net.delete_recordings()
net.cell(10).branch(0).loc(0.0).record()
```

--------------------------------

### JIT Compilation and Vmap for Efficient Simulation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Apply JIT compilation to the simulation function for speed and use vmap to parallelize simulations across multiple parameter sets. This significantly accelerates parameter sweeps.

```python
from jax import jit, vmap


cell = ...  # See tutorial on Basics of Jaxley.

def simulate(params):
    param_state = None
    param_state = cell.data_set("Na_gNa", params[0], param_state)
    param_state = cell.data_set("K_gK", params[1], param_state)
    return jx.integrate(cell, param_state=param_state, delta_t=0.025)

# Define 100 sets of sodium and potassium conductances.
all_params = jnp.asarray(np.random.rand(100, 2))

# Fast for-loops with jit compilation.
jitted_simulate = jit(simulate)
voltages = [jitted_simulate(params) for params in all_params]

# Using vmap for parallelization.
vmapped_simulate = vmap(jitted_simulate, in_axes=(0,))
voltages = vmapped_simulate(all_params)
```

--------------------------------

### Configure JAX Platform and Precision

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Set the JAX platform to 'cpu' and enable 64-bit precision. Use 'float32' for faster GPU performance, but be aware of potential stability issues with detailed neuron models.

```python
from jax import config
config.update("jax_platform_name", "cpu")
```

```python
config.update("jax_enable_x64", True)  # Set to false to use `float32`.
```

--------------------------------

### Run JIT Compiled Simulation (First Run)

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Execute the JIT-compiled simulation function. This initial run triggers the compilation process and will be slower.

```python
# First run, will be slow.
voltages = jitted_simulate(all_params[0])
```

--------------------------------

### Combine JIT and VMAP for Batched Parallel Simulations

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Combine `jit` and `vmap` to compile batches of parallel simulations, further optimizing performance for multiple simulation runs.

```python
jitted_vmapped_simulate = jit(vmap(simulate))
```

--------------------------------

### Import Jaxley and JAX Libraries

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Import necessary libraries for simulation, including JAX for numerical operations and Jaxley for building and simulating neural models.

```python
import matplotlib.pyplot as plt
import numpy as np
import jax.numpy as jnp
from jax import jit, vmap

import jaxley as jx
from jaxley.channels import Na, K, Leak
```

--------------------------------

### Initialize Fitting with Random Sample

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

Samples random initial parameters within defined bounds for the optimization process. Ensure 'upper_bounds' and 'lower_bounds' are defined prior to use.

```python
def sample_randomly():
    return jnp.asarray(np.random.rand(len(upper_bounds)) * (upper_bounds - lower_bounds) + lower_bounds)

_ = np.random.seed(0)
initial_params = sample_randomly()
```

--------------------------------

### Deploy MkDocs Documentation to GitHub

Source: https://github.com/jaxleyverse/jaxley/blob/main/mkdocs/README.md

Use this command to update the documentation on GitHub. It includes converting Jupyter notebooks to markdown before deployment.

```bash
jupyter nbconvert --to markdown ../docs/tutorials/*.ipynb --output-dir docs/tutorial/
mkdocs gh-deploy
```

--------------------------------

### Configure JAX for 64-bit Precision and CPU

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

Enables 64-bit floating-point precision and sets the JAX platform to CPU. Also configures memory fraction for XLA.

```python
from jax import config

config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import os
os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".8"
```

--------------------------------

### Training Loop with Optax

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Iterate through epochs, compute loss and gradients, update optimizer state, and apply parameter updates. The training progress is printed for each epoch.

```python
batch_size = 4
dataloader = Dataset(inputs, labels)
dataloader = dataloader.shuffle(seed=0).batch(batch_size)

for epoch in range(10):
    epoch_loss = 0.0
    
    for batch_ind, batch in enumerate(dataloader):
        current_batch, label_batch = batch
        loss_val, gradient = jitted_grad(opt_params, current_batch, label_batch)
        updates, opt_state = optimizer.update(gradient, opt_state)
        opt_params = optax.apply_updates(opt_params, updates)
        epoch_loss += loss_val
    
    print(f"epoch {epoch}, loss {epoch_loss}")
    
final_params = transform.forward(opt_params)
```

--------------------------------

### Order Imports with isort

Source: https://github.com/jaxleyverse/jaxley/blob/main/CONTRIBUTING.md

Run isort from the repository's root directory to consistently order imports.

```bash
isort
```

--------------------------------

### Define and Connect a Small Network

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/02_small_network.ipynb

This snippet demonstrates defining a network of cells, connecting them using `fully_connect`, and visualizing the network. It also shows how to modify synaptic parameters after connection.

```python
import jaxley as jx
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import fully_connect


# Define a network. `cell` is defined as in previous tutorial.
net = jx.Network([cell for _ in range(11)])

# Define synapses.
fully_connect(
    net.cell(range(10)),
    net.cell(10),
    IonotropicSynapse(),
)

# Change synaptic parameters.
net.select(edges=[0, 1]).set("IonotropicSynapse_gS", 0.1)  # nS

# Visualize the network.
net.compute_xyz()
fig, ax = plt.subplots(1, 1, figsize=(4, 4))
net.vis(ax=ax, detail="full", layers=[10, 1])  # or `detail="point"`.
```

--------------------------------

### Create and Connect a Network of Cells

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/08_importing_morphologies.ipynb

Load a morphology from an SWC file, create a network of identical cells, and establish synaptic connections between them using `IonotropicSynapse`.

```python
from jaxley.synapses import IonotropicSynapse

fname = "data/morph.swc"
cell = jx.read_swc(fname, ncomp=1)

net = jx.Network([cell]*5)
jx.connect(net.cell(0).soma.branch(0).comp(0), net[2,0,0], IonotropicSynapse())
jx.connect(net.cell(0).soma.branch(0).comp(0), net[3,0,0], IonotropicSynapse())
jx.connect(net.cell(0).soma.branch(0).comp(0), net[4,0,0], IonotropicSynapse())

jx.connect(net.cell(1).soma.branch(0).comp(0), net[2,0,0], IonotropicSynapse())
jx.connect(net.cell(1).soma.branch(0).comp(0), net[3,0,0], IonotropicSynapse())
jx.connect(net.cell(1).soma.branch(0).comp(0), net[4,0,0], IonotropicSynapse())

net.rotate(-90)

net.cell(0).move(0, 900)
net.cell(1).move(0, 1500)

net.cell(2).move(900, 600)
net.cell(3).move(900, 1200)
net.cell(4).move(900, 1800)

net.vis()
plt.axis("off")
plt.axis("square")
plt.show()
```

--------------------------------

### Setting Cell Parameters and Stimulating

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/12_simplified_models.ipynb

Configure cell parameters like length and radius, then apply external stimulation. Records the cell's voltage over time.

```python
cell.set("length", 1.0 / (2 * pi * 1e-5))
cell.set("radius", 1.0) # 1.0 is also the default.
cell.insert(Rate())
cell.stimulate(2.0 * jnp.ones((5,)))
cell.set("v", 2.0)
cell.record("v")

dt = 1.0
t_max = 10.0

v = jx.integrate(cell, t_max=t_max, delta_t=dt)
time_vec = jnp.arange(0, t_max + 2 * dt, dt)
fig, ax = plt.subplots(1, 1, figsize=(5, 2))
_ = plt.plot(time_vec, v.T, marker="o")
```

--------------------------------

### Initialize Izhikevich Neuron Model

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/12_simplified_models.ipynb

Sets up an Izhikevich neuron model by inserting the Izhikevich channel into a Jaxley Cell. Records the membrane voltage.

```python
from jaxley.channels import Izhikevich
cell = jx.Cell()
cell.insert(Izhikevich())
cell.record("v")
```

--------------------------------

### Integrate network with trainable parameters

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Runs the network simulation using the specified trainable parameters and simulation time.

```python
s = jx.integrate(net, params=params, t_max=5.0)
```

--------------------------------

### Setting up Checkpointing Levels

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

This code defines the checkpointing levels for memory optimization during simulation. It calculates the number of steps for each checkpoint based on the total simulation time.

```python
# For checkpointing.
checkpoint_levels = 2
checkpoints = [int(np.ceil(len(time_vec)**(1 / checkpoint_levels))) for _ in range(checkpoint_levels)]
```

--------------------------------

### Configure JAX for Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/00_jaxley_api.ipynb

Sets JAX to use 64-bit precision and the CPU, which can be important for numerical stability and compatibility in Jaxley simulations.

```python
from jax import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import connect
import matplotlib.pyplot as plt
import numpy as np
```

--------------------------------

### Jaxley Voltage-Clamp Experiment

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/05_channel_and_synapse_models.ipynb

Demonstrates how to perform a voltage-clamp experiment using the `.clamp()` method for debugging custom channel models. This involves setting up a cell, inserting a channel, recording its current, defining a clamped voltage protocol, and integrating the cell's response.

```python
cell = jx.Cell()
cell.insert(K())
cell.record("i_K")  # Record the `current_name`.

delay, dur, amp, baseline, dt, t_max = 10.0, 10.0, 100.0, -70.0, 0.025, 200.0
clamped_voltage = jx.step_current(delay, dur, amp, dt, t_max) + baseline

cell.clamp("v", clamped_voltage)
channel_current = jx.integrate(cell, delta_t=dt)
```

--------------------------------

### Visualize Network After Full Connection

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/02_small_network.ipynb

Renders the network visualization after applying `fully_connect`, showing the newly established synaptic connections.

```python
fig, ax = plt.subplots(1, 1, figsize=(3, 6))
_ = net.vis(ax=ax, detail="full")
```

--------------------------------

### Set up recordings for network output

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Configures specific locations within the network to record voltage, typically for output neurons.

```python
net.delete_recordings()
net.cell(0).branch(0).loc(0.0).record()
net.cell(1).branch(0).loc(0.0).record()
net.cell(2).branch(0).loc(0.0).record()
```

--------------------------------

### Integrate with Checkpointing

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Illustrates how to modify the integrate function call within the simulate function to enable gradient checkpointing.

```python
def simulate(params, inputs):
    # ...
    return jx.integrate(..., checkpoint_lengths=checkpoints)
```

--------------------------------

### Import necessary libraries

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Imports core Jax, Matplotlib, and Jaxley modules for network definition and simulation.

```python
from jax import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import matplotlib.pyplot as plt
import numpy as np
import jax
import jax.numpy as jnp
from jax import jit, vmap, value_and_grad

import jaxley as jx
from jaxley.channels import Leak
from jaxley.synapses import TanhRateSynapse
from jaxley.connect import fully_connect
```

--------------------------------

### Run Full Test Suite

Source: https://github.com/jaxleyverse/jaxley/wiki/Release-workflow

Execute the complete test suite for the Jaxley project. Ensure all tests pass before proceeding with a release.

```shell
pytest tests
```

--------------------------------

### Simulate Neuron with Step Current in Jaxley

Source: https://github.com/jaxleyverse/jaxley/blob/main/README.md

Demonstrates how to define a cell, insert Hodgkin-Huxley channels, stimulate with a step current, record voltage, and integrate the simulation. Set the JAX platform (CPU, GPU, or TPU) before simulation.

```python
import matplotlib.pyplot as plt
from jax import config

import jaxley as jx
from jaxley.channels import HH

config.update("jax_platform_name", "cpu")  # Or "gpu" / "tpu".

cell = jx.Cell()  # Define cell.
cell.insert(HH())  # Insert channels.

current = jx.step_current(i_delay=1.0, i_dur=1.0, i_amp=0.1, delta_t=0.025, t_max=10.0)
cell.stimulate(current)  # Stimulate with step current.
cell.record("v")  # Record voltage.

v = jx.integrate(cell)  # Run simulation.
plt.plot(v.T)  # Plot voltage trace.
```

--------------------------------

### Execute Batched Parallel Simulations

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Run simulations using the combined `jit` and `vmap` function within a loop to process multiple batches.

```python
for batch in range(10):
    all_params = jnp.asarray(np.random.rand(5, 2))
    voltages_batch = jitted_vmapped_simulate(all_params)
```

--------------------------------

### Run Gradient Descent with Polyak Learning Rate

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

Implements gradient descent using the Optax library with a Polyak-style learning rate. This optimizer adjusts the learning rate based on the current loss and gradient norm. Requires 'jaxley.optimize.utils.l2_norm' and 'optax'.

```python
import optax
from jaxley.optimize.utils import l2_norm

# We scale the learning rate with the number of parameters. Since
# the gradient gets divided by its own norm, it typically gets
# divided by larger numbers for models with many parameters.
# Choosing mu as below counteracts this.
#
# Also note that this learning rate is relatively high---you might
# want to explore smaller learning rates.
mu = 0.1 * l2_norm(jnp.ones(len(opt_params)))

beta = 0.9  # Good values are typically between 0.8 and 1.2.

optimizer = optax.inject_hyperparams(optax.sgd)(learning_rate=mu)
opt_state = optimizer.init(opt_params)

for epoch in range(10):
    loss_val, grad_val = grad_fn(opt_params)

    # Polyak style learning rate.
    grad_val = grad_val / l2_norm(grad_val)**beta
    opt_state.hyperparams["learning_rate"] = loss_val * mu
    
    # Update parameters and optimizer.
    updates, opt_state = optimizer.update(grad_val, opt_state)
    opt_params = optax.apply_updates(opt_params, updates)
    
    print(f"Loss in epoch {epoch}: {loss_val:.4f}")
```

--------------------------------

### Import Necessary Libraries

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

Imports core JAX, NumPy, Matplotlib, and JAXley modules for cell modeling and simulation.

```python
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpy as np

import jaxley as jx
from jaxley.channels import Leak
from jaxley_mech.channels.l5pc import *
from jaxley.morphology import distance_direct

# To suppress Pandas performance warnings:
import pandas as pd
import warnings
warnings.simplefilter("ignore", category=pd.errors.PerformanceWarning)
```

--------------------------------

### Import SWC file into networkX graph

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/dev/graph_backend.ipynb

Use `to_swc_graph` to load an SWC file into a networkX DiGraph. The directionality reflects the tracing direction.

```python
from jaxley.io.graph import to_swc_graph

swc_graph = to_swc_graph("../data/example.swc")
```

--------------------------------

### JIT and Vmapped Simulation Functions

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

This code prepares optimized versions of the simulation function using JAX. `jitted_sim` is a just-in-time compiled version for single simulations, while `vmapped_sim` is compiled and vectorized for batch simulations.

```python
jitted_sim = jax.jit(simulate)
vmapped_sim = jax.jit(jax.vmap(simulate, in_axes=(0,)))

```

--------------------------------

### Initialize Network and Connect Cells

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/10_advanced_parameter_sharing.ipynb

Initializes a network of six cells and creates fully connected ionotropic synapses between all cells. This sets up the basic network structure for parameter sharing experiments.

```python
from typing import Union, List

net = jx.Network([cell for _ in range(6)])
fully_connect(net.cell("all"), net.cell("all"), IonotropicSynapse())

net.cell([0, 1]).add_to_group("exc")
net.cell([2, 3]).add_to_group("inh")
net.cell([4, 5]).add_to_group("unknown")
```

--------------------------------

### Use Exponential Euler Solver

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/how_to_guide/choose_solver.ipynb

This snippet shows the basic usage of the 'exp_euler' solver within the `jx.integrate` function. It is generally stable and can provide speedups on GPUs.

```python
v = jx.integrate(..., solver="exp_euler")
```

--------------------------------

### Configure Parameter Sharing and Trainability

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/10_advanced_parameter_sharing.ipynb

Groups synapses by their pre- and post-synaptic cell types and assigns a unique parameter index to each group using `groupby().ngroup()`. It then makes the 'IonotropicSynapse_gS' parameter trainable for the selected synapses, resulting in 6 trainable parameters.

```python
# Step 6: use groupby to specify parameter sharing and make the parameters trainable.
subnetwork.edges["controlled_by_param"] = subnetwork.edges.groupby(["pre_cell_type", "post_cell_type"]).ngroup()
subnetwork.make_trainable("IonotropicSynapse_gS")
```

--------------------------------

### Build Package Distributions

Source: https://github.com/jaxleyverse/jaxley/wiki/Release-workflow

Use the build module to create source distribution (.tar.gz) and wheel (.whl) files for the Jaxley package. These are typically located in the 'dist/' directory.

```shell
python -m build
```

--------------------------------

### Import Jaxley and Supporting Libraries

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/01_morph_neurons.ipynb

Import necessary libraries for building and simulating neuron models, including Jaxley core components, channels, synapses, and plotting utilities.

```python
import matplotlib.pyplot as plt
import numpy as np
import jax.numpy as jnp
from jax import jit

import jaxley as jx
from jaxley.channels import Na, K, Leak
from jaxley.synapses import IonotropicSynapse
from jaxley.connect import fully_connect
```

--------------------------------

### Establish Regression Test Baseline Locally

Source: https://github.com/jaxleyverse/jaxley/wiki/Running-regression-tests

Run this command to establish a baseline for local regression tests on the main branch.

```shell
git checkout main
NEW_BASELINE=1 pytest -m regression
```

--------------------------------

### Setting Synaptic Conductance with Multiple Criteria

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/09_advanced_indexing.ipynb

Combine `copy_node_property_to_edges()`, `query()` with multiple conditions, and `select()` to set synaptic parameters based on both presynaptic and postsynaptic neuron indices.

```python
# Set synaptic conductance of all synapses that
# 1) have cells 2 or 3 as presynaptic neuron and
# 2) has cell 5 as postsynaptic neuron
df = net.edges
df = df.query("pre_global_cell_index in [2, 3]")
df = df.query("post_global_cell_index == 5")
net.select(edges=df.index).set("Ionotropic_gS", 0.3)
```

--------------------------------

### Initialize Plot for Branch Visualization

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/08_importing_morphologies.ipynb

Initializes a matplotlib figure and axes for visualizing individual branches of a morphology.

```python
fig, ax = plt.subplots(1, 1, figsize=(5, 5))
```

--------------------------------

### Integrate Jaxley Cell Simulation

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/13_graph_backend.ipynb

Runs the simulation for a configured Jaxley cell and returns the voltage trace.

```python
import jaxley as jx

v = jx.integrate(cell)
```

--------------------------------

### Customize Exponential Euler Solver for Performance

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Pre-compute the transition matrix for the exponential Euler solver to further optimize performance when simulating with consistent cell parameters.

```python
# You do not have to rerun `customize_solver_exp_euler` until you modify
# length, radius, axial resistivity, or capacitance.
cell.customize_solver_exp_euler(
    exp_euler_transition=cell.build_exp_euler_transition_matrix(delta_t)
)

v1 = jx.integrate(cell, solver="exp_euler")
cell.set("gLeak", 1e-4)
v2 = jx.integrate(cell, solver="exp_euler")
```

--------------------------------

### Run Vectorized Simulation with VMAP

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Execute the `vmap`-enabled simulation function on all parameter sets. This leverages GPU parallelization if available.

```python
voltages = vmapped_simulate(all_params)
```

--------------------------------

### JIT Compile a Simulation Function

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Use `jit` to compile a simulation function for faster subsequent runs. The first execution will be slower due to compilation.

```python
jitted_simulate = jit(simulate)
```

--------------------------------

### Import Jaxley and Core Libraries

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/05_channel_and_synapse_models.ipynb

Imports necessary libraries for Jaxley simulations, including configuration for 64-bit precision and CPU usage, plotting, numerical operations, and JAX functionalities.

```python
from jax import config
config.update("jax_enable_x64", True)
config.update("jax_platform_name", "cpu")

import matplotlib.pyplot as plt
import numpy as np
import jax
import jax.numpy as jnp
from jax import jit, value_and_grad

import jaxley as jx
```

--------------------------------

### Define Stimulation Protocol and Recording

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/examples/00_l5pc_gradient_descent.ipynb

Sets up a 90 ms step current stimulus of 1.8 nA applied to the soma and configures voltage recording.

```python
dt = 0.025  # ms
t_max = 100.0  # ms
time_vec = np.arange(0, t_max+2*dt, dt)

cell.delete_stimuli()
cell.delete_recordings()

i_delay = 5.0  # ms
i_dur = 90.0  # ms
i_amp = 1.8  # nA
current = jx.step_current(i_delay, i_dur, i_amp, dt, t_max)
cell.soma.branch(0).loc(0.5).stimulate(current)
cell.soma.branch(0).loc(0.5).record()

cell.set("v", -72.0)
cell.init_states()
```

--------------------------------

### Run JIT Compiled Simulation (Subsequent Runs)

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Run the JIT-compiled simulation function multiple times. These subsequent runs will be significantly faster after the initial compilation.

```python
# More runs, will be much faster.
voltages = jnp.asarray([jitted_simulate(params) for params in all_params])
print("voltages.shape", voltages.shape)
```

--------------------------------

### Visualize Specific Module Views

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/00_jaxley_api.ipynb

Create and visualize specific views of module components like cells, branches, and components within branches. This allows for detailed inspection of network parts.

```python
fig, ax = plt.subplots(1,1, figsize=(3,3))
net.cell(0).vis(ax=ax, color="blue") # View of the 0th cell of the network
net.cell(1).vis(ax=ax, color="red") # View of the 1st cell of the network

net.cell(0).branch(0).vis(ax=ax, color="green") # View of the 1st branch of the 0th cell of the network
net.cell(1).branch(1).comp(1).vis(ax=ax, color="black", type="line") # View of the 0th comp of the 1st branch of the 0th cell of the network
```

--------------------------------

### Precompute Transition Matrix for Exponential Euler

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/how_to_guide/choose_solver.ipynb

For significant speedups with the 'exp_euler' solver on GPUs, especially with repeated simulations using the same cell parameters, precompute the transition matrix once. This avoids redundant calculations within the integration loop.

```python
cell.customize_solver_exp_euler(
    exp_euler_transition=cell.build_exp_euler_transition_matrix(delta_t)
)
```

--------------------------------

### Integrate with Exponential Euler Solver on GPU

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Change the integration solver to 'exp_euler' for potential performance gains on GPU, especially with smaller morphologies.

```python
v = jx.integrate(cell, solver="exp_euler")
```

--------------------------------

### Prepare training data

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Generates synthetic input data and corresponding binary labels for a classification task.

```python
inputs = jnp.asarray(np.random.rand(100, 2))
labels = jnp.asarray((inputs[:, 0] + inputs[:, 1]) > 1.0)
```

--------------------------------

### Format Code with Black

Source: https://github.com/jaxleyverse/jaxley/blob/main/CONTRIBUTING.md

Run Black, an automatic Python code formatter, from the repository's root directory to format all files.

```bash
black .
```

--------------------------------

### Initialize Unit-Free Rate-Based Neuron Model

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/12_simplified_models.ipynb

Sets up a unit-free rate-based neuron model using the Rate channel. It configures cell dimensions to make external currents and synapses unit-free. Records membrane voltage.

```python
from jaxley.channels import Rate
from math import pi

cell = jx.Cell()
cell.set("length", 1.0 / (2 * pi * 1e-5))  # Make external current and synapses unit-free.
cell.insert(Rate())
cell.record("v")
```

--------------------------------

### Upload Package to PyPI

Source: https://github.com/jaxleyverse/jaxley/wiki/Release-workflow

Upload the built package distributions to the Python Package Index (PyPI) using twine. This command requires access to the Jaxley PyPI account.

```shell
twine upload dist/*
```

--------------------------------

### Vectorized Simulation Function with VMAP

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/04_jit_and_vmap.ipynb

Create a new simulation function using `vmap` to handle multiple parameter sets in parallel, suitable for GPU acceleration.

```python
# Using vmap for parallelization.
vmapped_simulate = vmap(jitted_simulate)
```

--------------------------------

### Stimulate and Record from Network

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/05_channel_and_synapse_models.ipynb

Applies a stimulus to a specific location in the network and sets up recordings for cells to monitor their voltage over time. This prepares the network for simulation.

```python
net.cell(0).branch(0).loc(0.0).stimulate(jx.step_current(1.0, 2.0, 0.1, 0.025, 10.0))
for i in range(3):
    net.cell(i).branch(0).loc(0.0).record()
```

--------------------------------

### Initialize Adam Optimizer

Source: https://github.com/jaxleyverse/jaxley/blob/main/docs/tutorials/07_gradient_descent.ipynb

Initialize the Adam optimizer with a specified learning rate. The optimizer state is initialized using the model's parameters.

```python
optimizer = optax.adam(learning_rate=0.01)
opt_state = optimizer.init(opt_params)
```