jax.numpy.shape#

jax.numpy.shape(a)[source]#

Return the shape an array.

JAX implementation ofnumpy.shape(). Unlikenp.shape, this functionraises aTypeError if the input is a collection such as a list ortuple.

Parameters:

a (ArrayLike |SupportsShape) – array-like object, or any object with ashape attribute.

Returns:

An tuple of integers representing the shape ofa.

Return type:

tuple[int, …]

Examples

Shape for arrays:

>>>x=jnp.arange(10)>>>jnp.shape(x)(10,)>>>y=jnp.ones((2,3))>>>jnp.shape(y)(2, 3)

This also works for scalars:

>>>jnp.shape(3.14)()

For arrays, this can also be accessed via thejax.Array.shape property:

>>>x.shape(10,)