code: def load_pretrained():
      return nnx.Linear(784, 128, rngs=nnx.Rngs(0))
  
imports: import jax
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling a global function with example usage: load_pretrained() and returns: nnx