code: def _assert_values_for_graph(self, padded_graph, wrapper):
      # Make padded graph values non zero.
      padded_graph = padded_graph._replace(
          nodes=tree.tree_map(lambda x: x - 1.0, padded_graph.nodes),
          edges=tree.tree_map(lambda x: x - 1.0, padded_graph.edges),
          globals=tree.tree_map(
              lambda x: x - 1.0, padded_graph.globals
          ),
      )
      true_valid_graph = utils.unbatch(padded_graph)[0]
      if wrapper:
          zeroed_graph_net = utils.with_zero_out_padding_outputs(
              lambda x: x
          )
          zeroed_padded_graph = zeroed_graph_net(padded_graph)
      else:
          zeroed_padded_graph = utils.zero_out_padding(padded_graph)
      graphs = utils.unbatch(zeroed_padded_graph)
      valid_graph = graphs[0]
      padding_graphs = graphs[1:]
      tree.tree_map(
          np.testing.assert_array_equal,
          valid_graph.nodes,
          true_valid_graph.nodes,
      )
      for padding_graph in padding_graphs:
          tree.tree_map(
              lambda x: np.testing.assert_array_equal(
                  x, jnp.zeros_like(x)
              ),
              padding_graph.nodes,
          )
  
imports: import functools
  import os
  from absl.testing import absltest
  from absl.testing import parameterized
  import jax
  from jax.lib import xla_bridge
  import jax.numpy
  import jax.tree_util
  from jraph._src import graph
  from jraph._src import utils
  import numpy
label_desc: Use image node_id 1 for calling the ZeroOutTest obj's underlying member method code with example usage: obj._assert_values_for_graph(padded_graph, wrapper) without return types