code: def get_2sat_problem(
      min_n_literals: int, max_n_literals: int
  ) -> Problem:
      """Creates bipartite-graph representing a randomly generated 2-sat problem.
  
      Args:
        min_n_literals: minimum number of literals in the 2-sat problem.
        max_n_literals: maximum number of literals in the 2-sat problem.
  
      Returns:
        bipartite-graph, node labels and node mask.
      """
      n_literals = random.randint(min_n_literals, max_n_literals)
      n_literals_true = random.randint(1, n_literals - 1)
      n_constraints = n_literals * (n_literals - 1) // 2
  
      n_node = n_literals + n_constraints
      # 0 indicates a literal node
      # 1 indicates a constraint node.
      nodes = [0 if i < n_literals else 1 for i in range(n_node)]
      edges = []
      senders = []
      for literal_node1 in range(n_literals):
          for literal_node2 in range(literal_node1 + 1, n_literals):
              senders.append(literal_node1)
              senders.append(literal_node2)
              # 1 indicates that the literal must be true for this constraint.
              # 0 indicates that the literal must be false for this constraint.
              # I.e. with literals a and b, we have the following possible constraints:
              # 0, 0 -> a or b
              # 1, 0 -> not a or b
              # 0, 1 -> a or not b
              # 1, 1 -> not a or not b
              edges.append(1 if literal_node1 < n_literals_true else 0)
              edges.append(1 if literal_node2 < n_literals_true else 0)
  
      graph = jraph.GraphsTuple(
          n_node=np.asarray([n_node]),
          n_edge=np.asarray([2 * n_constraints]),
          # One-hot encoding for nodes and edges.
          nodes=np.eye(2)[nodes],
          edges=np.eye(2)[edges],
          globals=None,
          senders=np.asarray(senders),
          receivers=np.repeat(np.arange(n_constraints) + n_literals, 2),
      )
  
      # In order to jit compile our code, we have to pad the nodes and edges of
      # the GraphsTuple to a static shape.
      max_n_constraints = max_n_literals * (max_n_literals - 1) // 2
      max_nodes = max_n_literals + max_n_constraints + 1
      max_edges = 2 * max_n_constraints
      graph = jraph.pad_with_graphs(graph, max_nodes, max_edges)
  
      # The ground truth solution for the 2-sat problem.
      labels = (np.arange(max_nodes) < n_literals_true).astype(np.int32)
      labels = np.eye(2)[labels]
  
      # For the loss calculation we create a mask for the nodes, which masks the
      # the constraint nodes and the padding nodes.
      mask = (np.arange(max_nodes) < n_literals).astype(np.int32)
      return Problem(graph=graph, labels=labels, mask=mask)
  
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 1 for calling a global function with example usage: get_2sat_problem(min_n_literals, max_n_literals) and returns: Problem