jax.numpy.diagonal#

jax.numpy.diagonal(a,offset=0,axis1=0,axis2=1)[source]#

Returns the specified diagonal of an array.

JAX implementation ofnumpy.diagonal().

The JAX version always returns a copy of the input, although if this is usedwithin a JIT compilation, the compiler may avoid the copy.

Parameters:

• a (ArrayLike) – Input array. Must be at least 2-dimensional.

• offset (int) – optional, default=0. Diagonal offset from the main diagonal.Must be a static integer value. Can be positive or negative.

• axis1 (int) – optional, default=0. The first axis along which to take the diagonal.

• axis2 (int) –

  optional, default=1. The second axis along which to take the diagonal.

  Returns:

  A 1D array for 2D input, and in general a N-1 dimensional arrayfor N-dimensional input.

Return type:

Array

See also

• jax.numpy.diag()

• jax.numpy.diagflat()

Examples

>>>x=jnp.array([[1,2,3],...[4,5,6],...[7,8,9]])>>>jnp.diagonal(x)Array([1, 5, 9], dtype=int32)>>>jnp.diagonal(x,offset=1)Array([2, 6], dtype=int32)>>>jnp.diagonal(x,offset=-1)Array([4, 8], dtype=int32)