code: def concatenated_args(
      update: Optional[Callable[..., ArrayTree]] = None,
      *,
      axis: int = -1,
  ) -> Union[
      Callable[..., ArrayTree],
      Callable[[Callable[..., ArrayTree]], ArrayTree],
  ]:
      """Decorator that concatenates arguments before being passed to an update_fn.
  
      By default node, edge and global features are passed separately to update
      functions. However, it is common practice to concatenate these features before
      passing them to a neural network. This wrapper concatenates the arguments
      for you.
  
      For example::
  
        # Without the wrapper
        def update_node_fn(nodes, receivers, globals):
          return net(jnp.concatenate([nodes, receivers, globals], axis=1))
  
        # With the wrapper
        @concatenated_args
        def update_node_fn(features):
          return net(features)
  
      Args:
        update: an update function that takes ``jnp.ndarray``.
        axis: the axis upon which to concatenate.
  
      Returns:
        A wrapped function with the arguments concatenated.
      """
  
      def _decorate(f):
          @functools.wraps(update)
          def wrapper(*args, **kwargs):
              combined_args = (
                  tree.tree_flatten(args)[0]
                  + tree.tree_flatten(kwargs)[0]
              )
              concat_args = jnp.concatenate(combined_args, axis=axis)
              return f(concat_args)
  
          return wrapper
  
      # If the update function is passed, then decorate the update function.
      if update:
          return _decorate(update)
  
      # Otherwise, return the decorator.
      return _decorate
  
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 27 for calling a global function with example usage: concatenated_args(update) and returns: _decorate, wrapper, _decorate, f