PyTorch
API: torch.trace
Strategy: Macro 'torch.sum(torch.diagonal({a}, {offset}, {axis1}, {axis2}))'
JAX (Core)
API: jnp.trace
Strategy: Direct Mapping
NumPy
API: np.trace
Strategy: Direct Mapping
Keras
API: keras.ops.trace
Strategy: Macro 'keras.ops.sum(keras.ops.diag({input}))'
TensorFlow
API: tf.linalg.trace
Strategy: Direct Mapping
Apple MLX
API: mlx.core.trace
Strategy: Macro 'mlx.core.sum(mlx.core.diagonal({input}))'
Flax NNX
API: jax.numpy.trace
Strategy: Direct Mapping
PaxML / Praxis
API: jax.numpy.trace
Strategy: Direct Mapping