code: def test_tree_at_replace_fn(getkey):
      key = getkey()
      pytree = [
          1,
          2,
          3,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key),
      ]
  
      def replace_fn(x):
          if isinstance(x, int):
              return "found an int"
          else:
              return x
  
      true_pytree1 = [
          "found an int",
          "found an int",
          3,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key),
      ]
      where = lambda tree: (tree[0], tree[1])
      pytree1 = eqx.tree_at(where, pytree, replace_fn=replace_fn)
  
      assert pytree1[:3] == true_pytree1[:3]
      assert jnp.all(pytree1[3]["a"] == true_pytree1[3]["a"])
      assert jnp.all(pytree1[-1].weight == true_pytree1[-1].weight)
      assert jnp.all(pytree1[-1].bias == true_pytree1[-1].bias)
  
      with pytest.raises(ValueError):
          eqx.tree_at(
              where, pytree, replace=(0, 1), replace_fn=replace_fn
          )
  
imports: import equinox
  import jax
  import jax.core
  import jax.nn
  import jax.numpy
  import jax.random
  import jax.tree_util
  import numpy
  import pytest
label_desc: Use image node_id 2 for calling a global function with example usage: test_tree_at_replace_fn(getkey) and returns: str, x