code: def get_predicted_logits(
      state: train_state.TrainState,
      graphs: jraph.GraphsTuple,
      rngs: Optional[Dict[str, jnp.ndarray]],
  ) -> jnp.ndarray:
      """Get predicted logits from the network for input graphs."""
      pred_graphs = state.apply_fn(state.params, graphs, rngs=rngs)
      logits = pred_graphs.globals
      return logits
  
imports: import os
  from typing import Any, Dict, Iterable, Tuple, Optional
  from absl import logging
  from clu import checkpoint
  from clu import metric_writers
  from clu import metrics
  from clu import parameter_overview
  from clu import periodic_actions
  import flax
  import flax.core
  import flax.linen
  from flax.training import train_state
  import jax
  import jax.numpy
  import jraph
  import ml_collections
  import numpy
  import optax
  import sklearn.metrics
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 7 for calling a global function with example usage: get_predicted_logits(state, graphs, rngs) and returns: logits