code: def make_embed_fn(latent_size):
      def embed(inputs):
          return nn.Dense(latent_size)(inputs)
  
      return embed
  
imports: import functools
  import logging
  import pathlib
  import pickle
  from typing import Sequence
  from absl import app
  from absl import flags
  from flax import linen
  from flax import optim
  import jax
  import jax.numpy
  import jraph
  from jraph.ogb_examples import data_utils
label_desc: Use image node_id 1 for calling a global function with example usage: make_embed_fn(latent_size) and returns: embed, nn