code: def test_single_arg(self, args_shapes, kwargs_shapes, axis):
      args = tree.tree_map(lambda x: np.random.randn(*x), args_shapes)
      kwargs = {
          k: np.random.randn(*shape)
          for k, shape in kwargs_shapes.items()
      }
  
      @utils.concatenated_args(axis=axis)
      def update_fn(feat):
          return feat
  
      out = update_fn(*args, **kwargs)
      expected_out = jnp.concatenate(
          list(tree.tree_flatten(args)[0])
          + list(tree.tree_flatten(kwargs)[0]),
          axis=axis,
      )
      np.testing.assert_allclose(out, expected_out)
  
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 ConcatenatedArgsWrapperTest obj's underlying member method code with example usage: obj.test_single_arg(args_shapes, kwargs_shapes, axis) and returns: feat