import torch
import jax
import numpy as np
import tensorflow as tf
import jax.numpy as jnp
A = np.random.rand(1000, 100, 100)
%timeit np.einsum("...ii", A) # 109 μs ± 1.71 μs per loop
%timeit np.trace(A, axis1=-2, axis2=-1) # 116 μs ± 1.79 μs
%timeit A.diagonal(axis1=-2, axis2=-1).sum(-1) # 114 μs ± 2.87 μs per loop
A = torch.rand(1000, 100, 100)
%timeit torch.einsum("...ii", A)  # 3.17 ms ± 1.1 ms per loop
%timeit A.diagonal(dim1=-2, dim2=-1).sum(-1) # 3.1 ms ± 879 μs per loop
A = tf.random.uniform((1000, 100, 100))
@tf.function
def trace_sum(A):
    return tf.einsum("...ii", A)
@tf.function
def trace_sum_v2(A):
    return tf.reduce_sum(tf.linalg.diag_part(A), axis=-1)
# Warm-up execution
trace_sum(A)
trace_sum_v2(A)
# Benchmarking
%timeit trace_sum(A) # 486 μs ± 21.1 μs per loop
%timeit trace_sum_v2(A) # 430 μs ± 36.1 μs per loop
%timeit tf.linalg.trace(A) # 404 μs ± 18.2 μs per loop
# For jax, the results might look different using jit
A = jnp.ones((1000, 100, 100))
%timeit jnp.einsum("...ii", A) # 13.6 ms ± 324 μs per loop
%timeit jax.vmap(jnp.trace)(A) # 12.1 ms ± 457 μs per loop
%timeit A.diagonal(axis1=-2, axis2=-1).sum(-1) # 1.64 ms ± 3.62 ms per loop