code: def test_intermediates_example_functional(self):
      class Linear(nnx.Module):
          def __init__(self, din: int, dout: int, *, rngs: nnx.Rngs):
              key = rngs.params()
              self.w = nnx.Param(jax.random.uniform(key, (din, dout)))
              self.b = nnx.Param(jnp.zeros((dout,)))
  
          def __call__(self, x):
              y = x @ self.w + self.b[None]
              self.y = nnx.Intermediate(y)
              return y
  
      model = Linear(12, 2, rngs=nnx.Rngs(0))
  
      state, graphdef = model.split()
  
      y, (state, _) = graphdef.apply(state)(jnp.ones((8, 12)))
  
      intermediates, state = state.split(nnx.Intermediate, ...)
  
      assert "y" in intermediates
  
file_path: flax/flax/experimental/nnx/tests/test_integration.py
imports: import typing
  import jax
  import jax.numpy
  import numpy
  from flax.experimental import nnx
label_desc: Use image node_id 6 for calling the TestIntegration obj's underlying member method code with example usage: obj.test_intermediates_example_functional() and returns: y