code: def test_stateful(tmp_path):
      class Model(eqx.Module):
          norm1: eqx.nn.BatchNorm
          norm2: eqx.nn.BatchNorm
  
      model = Model(
          eqx.nn.BatchNorm(3, "hi"), eqx.nn.BatchNorm(4, "bye")
      )
      state = eqx.nn.State(model)
  
      eqx.tree_serialise_leaves(tmp_path, (model, state))
  
      model2 = Model(
          eqx.nn.BatchNorm(3, "hi"), eqx.nn.BatchNorm(4, "bye")
      )
      state2 = eqx.nn.State(model2)
  
      eqx.tree_deserialise_leaves(tmp_path, (model2, state2))
  
file_path: equinox/tests/test_serialisation.py
imports: import os
  import equinox
  import jax
  import jax.numpy
  import numpy
  import pytest
  from jax.dtypes import bfloat16
label_desc: Use image node_id 8 for calling a global function with example usage: test_stateful(tmp_path) without return types