code: def test_helpful_errors(getkey, tmp_path):
      # Test that we get a helpful error message when the loading itself fails
      tree = jnp.array(1), {}
      eqx.tree_serialise_leaves(tmp_path, tree)
      bad_like_tree = (
          jnp.array(2),
          {
              "a": jnp.array(2),
              "b": jnp.array(2),
          },
      )
      with pytest.raises(
          RuntimeError,
          match=r"Error at leaf with path \(SequenceKey\(idx=1\), DictKey\(key='a'\)\)",
      ):
          _ = eqx.tree_deserialise_leaves(tmp_path, bad_like_tree)
  
      # Test that we get a helpful error message when the types don't match
      tree = (
          jnp.array(1),
          {
              "a": jnp.array(1),
              "b": jnp.array(1),
          },
      )
      eqx.tree_serialise_leaves(tmp_path, tree)
      bad_like_tree = (
          jnp.array(2),
          {
              "a": jnp.array(2, dtype=jnp.float32),
              "b": jnp.array(2),
          },
      )
      with pytest.raises(
          RuntimeError,
          match=r"Deserialised leaf at path \(SequenceKey\(idx=1\), DictKey\(key='a'\)\)",
      ):
          _ = eqx.tree_deserialise_leaves(tmp_path, bad_like_tree)
  
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 4 for calling a global function with example usage: test_helpful_errors(getkey, tmp_path) without return types