code: def train(
      data_path,
      master_csv_path,
      split_path,
      batch_size,
      num_training_steps,
      save_dir,
  ):
      """OGB Training Script."""
  
      # Initialize the dataset reader.
      reader = data_utils.DataReader(
          data_path=data_path,
          master_csv_path=master_csv_path,
          split_path=split_path,
          batch_size=batch_size,
      )
      # Repeat the dataset forever for training.
      reader.repeat()
  
      net = GraphNetwork(mlp_features=(128, 128), latent_size=128)
  
      # Get a candidate graph and label to initialize the network.
      graph = reader.get_graph_by_idx(0)
  
      # Initialize the network.
      logging.info("Initializing network.")
      params = net.init(jax.random.PRNGKey(42), graph)
      optimizer = optim.Adam(learning_rate=1e-4).create(params)
      optimizer = jax.device_put(optimizer)
  
      for idx in range(num_training_steps):
          graph = next(reader)
          # Jax will re-jit your graphnet every time a new graph shape is encountered.
          # In the limit, this means a new compilation every training step, which
          # will result in *extremely* slow training. To prevent this, pad each
          # batch of graphs to the nearest power of two. Since jax maintains a cache
          # of compiled programs, the compilation cost is amortized.
          graph = pad_graph_to_nearest_power_of_two(graph)
  
          # Remove the label from the input graph/
          label = graph.globals["label"]
          graph = graph._replace(globals={})
  
          optimizer, scalars = train_step(optimizer, graph, label, net)
          if idx % 100 == 0:
              logging.info(
                  "step: %s, loss: %s, acc: %s",
                  idx,
                  scalars["loss"],
                  scalars["accuracy"],
              )
      if save_dir is not None:
          with pathlib.Path(save_dir, "molhiv.pkl").open("wb") as fp:
              logging.info("Saving model to %s", save_dir)
              pickle.dump(optimizer.target, fp)
      logging.info("Training finished")
  
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 7 for calling a global function with example usage: train(data_path, master_csv_path, split_path, batch_size, num_training_steps, save_dir) without return types