code: def _get_deep_sets(graphs_tuple):
      node_fn = lambda n, g: jnp.concatenate((n, g), axis=-1)
      global_fn = lambda n: n * 2
      out = models.DeepSets(node_fn, global_fn)(graphs_tuple)
      num_graphs = len(graphs_tuple.n_node)
      num_nodes = len(graphs_tuple.nodes)
      broadcasted_globals = jnp.repeat(
          graphs_tuple.globals,
          graphs_tuple.n_node,
          total_repeat_length=num_nodes,
          axis=0,
      )
      expected_nodes = jnp.concatenate(
          (graphs_tuple.nodes, broadcasted_globals), axis=-1
      )
      node_gr_idx = jnp.repeat(
          jnp.arange(num_graphs),
          graphs_tuple.n_node,
          total_repeat_length=num_nodes,
      )
      expected_out = graphs_tuple._replace(
          nodes=expected_nodes,
          globals=utils.segment_sum(
              expected_nodes, node_gr_idx, num_segments=num_graphs
          )
          * 2,
      )
      return out, expected_out
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import models
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 12 for calling a global function with example usage: _get_deep_sets(graphs_tuple) and returns: out,  expected_out