code: def trace(
      x: JaxArray,
      /,
      *,
      offset: int = 0,
      axis1: int = 0,
      axis2: int = 1,
      out: Optional[JaxArray] = None,
  ) -> JaxArray:
      return jnp.trace(
          x, offset=offset, axis1=axis1, axis2=axis2, out=out
      )
  
imports: from collections import namedtuple
  from typing import Union, Optional, Tuple, Literal, Sequence, NamedTuple
  import jax.numpy
  import ivy
  from ivy import inf
  from ivy.func_wrapper import with_unsupported_dtypes
  from ivy.functional.backends.jax import JaxArray
  from .None import backend_version
  from ivy import promote_types_of_inputs
label_desc: Use image node_id 24 for calling a global function with example usage: trace() and returns: jnp