Trace

Returns the sum of the elements of the diagonal of the input 2-D matrix.

Abstract Signature:

Trace(input: Tensor)

PyTorch

API: torch.trace
Strategy: Macro 'torch.sum(torch.diagonal({a}, {offset}, {axis1}, {axis2}))'

JAX (Core)

API: jnp.trace
Strategy: Direct Mapping
Official Docs ↗

NumPy

API: np.trace
Strategy: Direct Mapping
Official Docs ↗

Keras

API: keras.ops.trace
Strategy: Macro 'keras.ops.sum(keras.ops.diag({input}))'

TensorFlow

API: tf.linalg.trace
Strategy: Direct Mapping
Official Docs ↗

Apple MLX

API: mlx.core.trace
Strategy: Macro 'mlx.core.sum(mlx.core.diagonal({input}))'

Flax NNX

API: jax.numpy.trace
Strategy: Direct Mapping
Official Docs ↗

PaxML / Praxis

API: jax.numpy.trace
Strategy: Direct Mapping
Official Docs ↗