code: def test_tree_at_replace(getkey):
      key = getkey()
      key1, key2 = jrandom.split(key, 2)
      pytree = [
          1,
          2,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key1),
      ]
      true_pytree1 = [1, 2, {"a": "hi"}, eqx.nn.Linear(1, 2, key=key1)]
      true_pytree2 = [
          1,
          2,
          {"a": jnp.array([1.0, 2.0])},
          eqx.nn.Linear(1, 2, key=key2),
      ]
      where1 = lambda tree: tree[2]["a"]
      where2 = lambda tree: (tree[3].weight, tree[3].bias)
      weight2 = true_pytree2[3].weight
      bias2 = true_pytree2[3].bias
      pytree1 = eqx.tree_at(where1, pytree, replace="hi")
      pytree2 = eqx.tree_at(where2, pytree, replace=(weight2, bias2))
  
      assert pytree1[:-2] == true_pytree1[:-2]
      assert pytree2[:-2] == true_pytree2[:-2]
      assert jnp.all(pytree1[-2]["a"] == true_pytree1[-2]["a"])
      assert jnp.all(pytree2[-2]["a"] == true_pytree2[-2]["a"])
      assert jnp.all(pytree1[-1].weight == true_pytree1[-1].weight)
      assert jnp.all(pytree1[-1].bias == true_pytree1[-1].bias)
      assert jnp.all(pytree2[-1].weight == true_pytree2[-1].weight)
      assert jnp.all(pytree2[-1].bias == true_pytree2[-1].bias)
  
      true_pytree3 = ["hi", 2, {"a": 4}, eqx.nn.Linear(1, 2, key=key1)]
      where3 = lambda tree: (tree[0], tree[2]["a"])
      pytree3 = eqx.tree_at(where3, pytree, replace=("hi", 4))
      assert pytree3[:-1] == true_pytree3[:-1]
      assert jnp.all(pytree3[-1].weight == true_pytree3[-1].weight)
      assert jnp.all(pytree3[-1].bias == true_pytree3[-1].bias)
  
      with pytest.raises(TypeError):
          eqx.tree_at(where3, pytree, replace=4)
      with pytest.raises(ValueError):
          eqx.tree_at(where3, pytree, replace=(3, 4, 5))
  
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 1 for calling a global function with example usage: test_tree_at_replace(getkey) without return types