code: def trace(x, /, *, offset=0, dtype=None):
      """
      Returns the sum along the specified diagonals of a matrix (or a stack of matrices) x.
      """
      x = _promote_to_default_dtype(x)
      return jax.numpy.trace(
          x, offset=offset, dtype=dtype, axis1=-2, axis2=-1
      )
  
imports: from functools import partial
  from typing import NamedTuple
  import jax
  from jax.experimental.array_api._data_type_functions import _promote_to_default_dtype
label_desc: Use image node_id 21 for calling a global function with example usage: trace() and returns: jax