code: def get_valid_mask(
      labels: jnp.ndarray, graphs: jraph.GraphsTuple
  ) -> jnp.ndarray:
      """Gets the binary mask indicating only valid labels and graphs."""
      # We have to ignore all NaN values - which indicate labels for which
      # the current graphs have no label.
      labels_mask = ~jnp.isnan(labels)
  
      # Since we have extra 'dummy' graphs in our batch due to padding, we want
      # to mask out any loss associated with the dummy graphs.
      # Since we padded with `pad_with_graphs` we can recover the mask by using
      # get_graph_padding_mask.
      graph_mask = jraph.get_graph_padding_mask(graphs)
  
      # Combine the mask over labels with the mask over graphs.
      return labels_mask & graph_mask[:, None]
  
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 8 for calling a global function with example usage: get_valid_mask(labels, graphs) and returns: unknown