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