code: def test_tree_at_subtree(getkey):
      class L(eqx.Module):
          def __call__(self, x):
              return x
  
      mlp = eqx.nn.MLP(2, 2, 2, 2, key=getkey())
  
      # m.layers is a node in the PyTree
      newmlp1 = eqx.tree_at(
          lambda m: m.layers, mlp, [L() for _ in range(len(mlp.layers))]
      )
  
      # tuple(m.layers) is a sequence of nodes in the PyTree.
      newmlp2 = eqx.tree_at(
          lambda m: tuple(m.layers),
          mlp,
          [L() for _ in range(len(mlp.layers))],
      )
  
      x = jrandom.normal(getkey(), (2,))
      assert (jnn.relu(x) == newmlp1(x)).all()
      assert (jnn.relu(x) == newmlp2(x)).all()
  
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 3 for calling a global function with example usage: test_tree_at_subtree(getkey) and returns: x