<commit_before>---
jupytext:
  formats: ipynb,md:myst
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.13.8
kernelspec:
  display_name: Python 3
  name: python3
---

+++ {"id": "8Z1q_WRdTpgv"}

[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/state_params.ipynb)
[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/notebooks/state_params.ipynb)

# Managing Parameters and State

+++ {"id": "Z7Qjr14xTsp7"}

We will show you how to…

- manage the variables from initialization to updates.
- split and re-assemble parameters and state.
- use vmap with batch-dependant state.

```{code-cell} ipython3
:id: KL1SWMpnUTAA
:tags: [skip-execution]

#@title
! pip install --upgrade pip jax jaxlib
! pip install --upgrade git+https://github.com/google/flax.git
! pip install optax
```

```{code-cell} ipython3
from functools import partial

import jax
from jax import numpy as jnp
from jax import random
import optax

from flax import linen as nn


#@title
# Initialize random variables
dummy_input = jnp.ones((32, 5))

X = random.uniform(random.PRNGKey(0), (128, 5),  minval=0.0, maxval=1.0)
noise = random.uniform(random.PRNGKey(0), (),  minval=0.0, maxval=0.1)
X += noise

W = random.uniform(random.PRNGKey(0), (5, 1),  minval=0.0, maxval=1.0)
b = random.uniform(random.PRNGKey(0), (),  minval=0.0, maxval=1.0)

Y = jnp.matmul(X, W) + b

num_epochs = 5
```

```{code-cell} ipython3
:id: eO6VTPVVT702

class BiasAdderWithRunningMean(nn.Module):
  momentum: float = 0.9

  @nn.compact
  def __call__(self, x):
    is_initialized = self.has_variable('batch_stats', 'mean')
    mean = self.variable('batch_stats', 'mean', jnp.zeros, x.shape[1:])
    bias = self.param('bias', lambda rng, shape: jnp.zeros(shape), x.shape[1:])
    if is_initialized:
      mean.value = (self.momentum * mean.value +
                    (1.0 - self.momentum) * jnp.mean(x, axis=0, keepdims=True))
    return mean.value + bias
```

+++ {"id": "bJ3TYyIaTy8f"}

This example model is a minimal example that contains both parameters (declared with `self.param`) and state variables (declared with `self.variable`).

The tricky part with initialization here is that we need to split the state variables and the parameters we’re going to optimize for.

First we define `update_step` as follow (with dummy loss that should be replaced for yours):

```{code-cell} ipython3
:id: B7xxRU02VeJb

@partial(jax.jit, static_argnums=(0,))
def update_step(apply_fn, x, opt_state, params, state):
  def loss(params):
    y, updated_state = apply_fn({'params': params, **state},
                                x, mutable=list(state.keys()))
    l = ((x - y) ** 2).sum() # Replace with your loss here.
    return l, updated_state

  (l, updated_state), grads = jax.value_and_grad(
      loss, has_aux=True)(params)
  updates, opt_state = tx.update(grads, opt_state)  # Defined below.
  params = optax.apply_updates(params, updates)
  return opt_state, params, state
```

+++ {"id": "zM_EyHqwVlEw"}

Then we can write the actual training code.

```{code-cell} ipython3
:id: 8RUFi57GVktj

model = BiasAdderWithRunningMean()
variables = model.init(random.PRNGKey(0), dummy_input)
# Split state and params (which are updated by optimizer).<commit_msg>Update ``state_params.md`` alongside ``state_params.ipynb``<commit_after>---
jupytext:
  formats: ipynb,md:myst
  text_representation:
    extension: .md
    format_name: myst
    format_version: 0.13
    jupytext_version: 1.13.8
kernelspec:
  display_name: Python 3
  name: python3
---

+++ {"id": "8Z1q_WRdTpgv"}

[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/flax/blob/main/docs/notebooks/state_params.ipynb)
[![Open On GitHub](https://img.shields.io/badge/Open-on%20GitHub-blue?logo=GitHub)](https://github.com/google/flax/blob/main/docs/notebooks/state_params.ipynb)

# Managing Parameters and State

+++ {"id": "Z7Qjr14xTsp7"}

We will show you how to…

- manage the variables from initialization to updates.
- split and re-assemble parameters and state.
- use vmap with batch-dependant state.

```{code-cell} ipython3
:id: KL1SWMpnUTAA
:tags: [skip-execution]

#@title
! pip install --upgrade pip jax jaxlib
! pip install --upgrade git+https://github.com/google/flax.git
! pip install optax
```

```{code-cell} ipython3
from functools import partial

import jax
from jax import numpy as jnp
from jax import random
import optax

from flax import linen as nn


#@title
# Initialize random variables
dummy_input = jnp.ones((32, 5))

X = random.uniform(random.PRNGKey(0), (128, 5),  minval=0.0, maxval=1.0)
noise = random.uniform(random.PRNGKey(0), (),  minval=0.0, maxval=0.1)
X += noise

W = random.uniform(random.PRNGKey(0), (5, 1),  minval=0.0, maxval=1.0)
b = random.uniform(random.PRNGKey(0), (),  minval=0.0, maxval=1.0)

Y = jnp.matmul(X, W) + b

num_epochs = 5
```

```{code-cell} ipython3
:id: eO6VTPVVT702

class BiasAdderWithRunningMean(nn.Module):
  momentum: float = 0.9

  @nn.compact
  def __call__(self, x):
    is_initialized = self.has_variable('batch_stats', 'mean')
    mean = self.variable('batch_stats', 'mean', jnp.zeros, x.shape[1:])
    bias = self.param('bias', lambda rng, shape: jnp.zeros(shape), x.shape[1:])
    if is_initialized:
      mean.value = (self.momentum * mean.value +
                    (1.0 - self.momentum) * jnp.mean(x, axis=0, keepdims=True))
    return mean.value + bias
```

+++ {"id": "bJ3TYyIaTy8f"}

This example model is a minimal example that contains both parameters (declared with `self.param`) and state variables (declared with `self.variable`).

The tricky part with initialization here is that we need to split the state variables and the parameters we’re going to optimize for.

First we define `update_step` as follow (with dummy loss that should be replaced for yours):

```{code-cell} ipython3
:id: B7xxRU02VeJb

@partial(jax.jit, static_argnums=(0,))
def update_step(apply_fn, x, opt_state, params, state):
  def loss(params):
    y, updated_state = apply_fn({'params': params, **state},
                                x, mutable=list(state.keys()))
    l = ((x - y) ** 2).sum() # Replace with your loss here.
    return l, updated_state

  (l, updated_state), grads = jax.value_and_grad(
      loss, has_aux=True)(params)
  updates, opt_state = tx.update(grads, opt_state)  # Defined below.
  params = optax.apply_updates(params, updates)
  return opt_state, params, updated_state
```

+++ {"id": "zM_EyHqwVlEw"}

Then we can write the actual training code.

```{code-cell} ipython3
:id: 8RUFi57GVktj

model = BiasAdderWithRunningMean()
variables = model.init(random.PRNGKey(0), dummy_input)
# Split state and params (which are updated by optimizer).
{'fraction_non_alphanumeric': 0.14310394471638352, 'fraction_numerical': 0.021595162683558884, 'mean_word_length': 4.168898809523809, 'pattern_counts': {'":': 8, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 10, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4415256', 'n_tokens_mistral': 2654, 'n_tokens_neox': 2531, 'n_words': 640}