jax.numpy.square#
- jax.numpy.square(x, /)[source]#
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