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()
  
      # Transform impure `net_fn` to pure functions with hk.transform.
      net = hk.without_apply_rng(hk.transform(net_fn))
      # 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)
      # Initialize the optimizer.
      opt_init, opt_update = optax.adam(1e-4)
      opt_state = opt_init(params)
  
      compute_loss_fn = functools.partial(compute_loss, net=net)
      # We jit the computation of our loss, since this is the main computation.
      # Using jax.jit means that we will use a single accelerator. If you want
      # to use more than 1 accelerator, use jax.pmap. More information can be
      # found in the jax documentation.
      compute_loss_fn = jax.jit(
          jax.value_and_grad(compute_loss_fn, has_aux=True)
      )
  
      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)
  
          # Extract the label from the graph.
          label = graph.globals["label"]
          graph = graph._replace(globals={})
  
          (loss, acc), grad = compute_loss_fn(params, graph, label)
          updates, opt_state = opt_update(grad, opt_state, params)
          params = optax.apply_updates(params, updates)
          if idx % 100 == 0:
              logging.info(
                  "step: %s, loss: %s, acc: %s", idx, loss, acc
              )
      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(params, fp)
      logging.info("Training finished")
  
imports: import functools
  import logging
  import pathlib
  import pickle
  from absl import app
  from absl import flags
  import haiku
  import jax
  import jax.numpy
  import jraph
  from jraph.ogb_examples import data_utils
  import optax
label_desc: Use image node_id 8 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