code: def pack_nodes_and_edges(
      batch_size: int,
      neighbor_config: configs.GraphRegConfig,
      features: Dict[Text, tf.Tensor],
      *args,
  ):
      """Pack features into a graph."""
      if neighbor_config.max_neighbors == 0:
          graph = {key: value for key, value in features.items()}
          graph["n_node"] = tf.ones(batch_size, dtype=tf.int64)
          graph["n_edge"] = tf.zeros(batch_size, dtype=tf.int64)
          graph["edges"] = tf.zeros([0, 1], dtype=tf.float32)
          graph["senders"] = graph["receivers"] = tf.zeros(
              0, dtype=tf.int64
          )
          graph["globals"] = tf.zeros([batch_size, 0], tf.float32)
          return (graph,) + args if args else graph
      ragged_fn = functools.partial(
          _make_ragged, batch_size, neighbor_config.max_neighbors
      )
      (
          features,
          neighbor_features,
          neighbor_weights,
      ) = keras.layers.NeighborFeatures(neighbor_config)(features)
      neighbor_weights = ragged_fn(tf.squeeze(neighbor_weights))
      neighbor_weights.flat_values.set_shape(
          [batch_size * neighbor_config.max_neighbors]
      )
      mask = neighbor_weights.with_flat_values(
          neighbor_weights.flat_values > 0
      )
      neighbor_features = tf.nest.map_structure(
          lambda x: tf.ragged.boolean_mask(ragged_fn(x), mask),
          neighbor_features,
      )
      neighbor_weights = tf.ragged.boolean_mask(neighbor_weights, mask)
      graph = {}
      for key, value in features.items():
          graph[key] = tf.concat(
              [
                  neighbor_features[key],
                  tf.RaggedTensor.from_row_splits(
                      value, tf.range(batch_size + 1, dtype=tf.int64)
                  ),
              ],
              1,
          )
      # Just take it from the last feature and assume it's all the same.
      center_node_ids = list(graph.values())[-1].row_limits() - 1
      n_node = list(graph.values())[-1].row_lengths()
      graph["edges"] = tf.reshape(neighbor_weights.flat_values, [-1, 1])
      graph["senders"] = (
          tf.range(
              tf.shape(neighbor_weights.flat_values)[0], dtype=tf.int64
          )
          + neighbor_weights.value_rowids()
      )
      graph["receivers"] = tf.gather(
          center_node_ids, neighbor_weights.value_rowids()
      )
      graph["n_edge"] = neighbor_weights.row_lengths()
      graph["n_node"] = n_node
      graph["globals"] = tf.zeros([batch_size, 0], tf.float32)
      graph = {
          key: value.flat_values
          if isinstance(value, tf.RaggedTensor)
          else value
          for key, value in graph.items()
      }
      return (graph,) + args if args else graph
  
imports: import functools
  from typing import Dict, Optional, Text
  import graph_nets
  import neural_structured_learning.configs
  import neural_structured_learning.keras
  import neural_structured_learning.lib
  import sonnet.v2
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: pack_nodes_and_edges(batch_size, neighbor_config, features) and returns: unknown, unknown