code: def train(num_steps: int):
      """Trains a graph neural network on a 2-sat problem."""
      train_dataset = (2, 15)
      test_dataset = (16, 20)
      random.seed(42)
  
      network = hk.without_apply_rng(hk.transform(network_definition))
      problem = get_2sat_problem(*train_dataset)
      params = network.init(jax.random.PRNGKey(42), problem.graph)
  
      @jax.jit
      def prediction_loss(params, problem):
          decoded_nodes = network.apply(params, problem.graph)
          # We interpret the decoded nodes as a pair of logits for each node.
          log_prob = jax.nn.log_softmax(decoded_nodes) * problem.labels
          return -jnp.sum(log_prob * problem.mask[:, None]) / jnp.sum(
              problem.mask
          )
  
      opt_init, opt_update = optax.adam(2e-4)
      opt_state = opt_init(params)
  
      @jax.jit
      def update(params, opt_state, problem):
          g = jax.grad(prediction_loss)(params, problem)
          updates, opt_state = opt_update(g, opt_state)
          return optax.apply_updates(params, updates), opt_state
  
      for step in range(num_steps):
          problem = get_2sat_problem(*train_dataset)
          params, opt_state = update(params, opt_state, problem)
          if step % 1000 == 0:
              train_loss = jnp.mean(
                  jnp.asarray(
                      [
                          prediction_loss(
                              params, get_2sat_problem(*train_dataset)
                          )
                          for _ in range(100)
                      ]
                  )
              ).item()
              test_loss = jnp.mean(
                  jnp.asarray(
                      [
                          prediction_loss(
                              params, get_2sat_problem(*test_dataset)
                          )
                          for _ in range(100)
                      ]
                  )
              ).item()
              logging.info(
                  "step %r loss train %r test %r",
                  step,
                  train_loss,
                  test_loss,
              )
  
file_path: jraph/jraph/examples/sat.py
imports: import collections
  import logging
  import random
  from absl import app
  import haiku
  import jax
  import jax.numpy
  import jraph
  import numpy
  import optax
label_desc: Use image node_id 3 for calling a global function with example usage: train(num_steps) and returns: unknown, optax,  opt_state