jax.numpy.square

jax.numpy.square(x, /)

Return the element-wise square of the input.

LAX-backend implementation of numpy.square().

Original docstring below.

Parameters:

x (array_like) – Input data.

Returns:

out – Element-wise x*x, of the same shape and dtype as x. This is a scalar if x is a scalar.

Return type:

ndarray or scalar