jax.Array.trace#

abstractArray.trace(offset=0,axis1=0,axis2=1,dtype=None,out=None)[source]#

Return the sum along the diagonal.

Refer tojax.numpy.trace() for full documentation.

Parameters:

• self (Array)

• offset (int |ArrayLike)

• axis1 (int)

• axis2 (int)

• dtype (DTypeLike |None)

• out (None)

Return type:

Array