code: def zero_out_padding(
      graph: gn_graph.GraphsTuple,
  ) -> gn_graph.GraphsTuple:
      """Zeroes out padded graphs values.
  
      Padded graphs can cause numeric overflow issues when a node has a large number
      connecting edges, and the graph network has many layers. By inserting this
      zeroing function between layers of the graph network these issues may be
      addressed.
  
      Zeroing is performed by multiplying a boolean mask with the corresponding
      array of the input dtype, assuming correct type casting happens.
  
      For example::
  
        # Set up padded jraph.GraphsTuple & simple jraph.GraphNetwork.
        key = jax.random.PRNGKey(0)
        graph = jraph.GraphsTuple(nodes=jax.random.randint(key, (3,), 0, 1,
        dtype=jnp.int32),
                          edges=jax.random.randint(key, (5,), 0, 1,
                          dtype=jnp.int32),
                          senders=jnp.array([0,0,1,1,2]),
                          receivers=[1,2,0,2,1],
                          n_node=jnp.array([3]),
                          n_edge=jnp.array([5]),
                          globals=jax.random.randint(key, (1,), 0, 128,
                          dtype=jnp.int32))
        padded_graph = jraph.pad_with_graphs(graph, n_node=4, n_edge=100, n_graph=2)
  
        def net_fn(graph: jraph.GraphsTuple) -> jraph.GraphsTuple:
          embedder = jraph.GraphMapFeatures(
              hk.Embed(2, 1), hk.Embed(2, 1), hk.Embed(2, 1))
          update_fn = jraph.concatenated_args(lambda x: jnp.sum(x, -1,
          keepdims=True))
          net = jraph.GraphNetwork(
              update_node_fn=update_fn,
              update_edge_fn=update_fn,
              update_global_fn=update_fn)
          embedded = embedder(graph)
          return net(embedded)
  
        net = haiku.without_apply_rng(haiku.transform(net_fn))
        params = net.init(key, padded_graph)
  
        # Without zeroing.
        net.apply(params, padded_graph).nodes
        >> [[ -1.1949954]
            [ -1.4456191]
            [ -1.1949954]
            [-88.505775 ]] # Very large activation in a single node.
  
        # With zeroing.
        zero_out_padding(net.apply(params, padded_graph)).nodes
        >> [[ -1.1949954]
            [ -1.4456191]
            [ -1.1949954]
            [-0. ]]  # Zeroed out activation.
  
      Args:
        graph: A padded graph.
  
      Returns:
        A graph with the same valid values as input, but padded values zeroed out.
      """
      edge_mask = get_edge_padding_mask(graph)
      masked_edges = jax.tree_map(_get_zero_fn(edge_mask), graph.edges)
      node_mask = get_node_padding_mask(graph)
      masked_nodes = jax.tree_map(_get_zero_fn(node_mask), graph.nodes)
      global_mask = get_graph_padding_mask(graph)
      masked_globals = jax.tree_map(
          _get_zero_fn(global_mask), graph.globals
      )
      return graph._replace(
          nodes=masked_nodes, edges=masked_edges, globals=masked_globals
      )
  
imports: import functools
  from typing import Any, Callable, Generator, Iterable, Iterator, List, Mapping, Optional, Sequence, Union
  import jax
  from jax import lax
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  import numpy
label_desc: Use image node_id 36 for calling a global function with example usage: zero_out_padding(graph) and returns: graph