jax.lax.square# jax.lax.square(x)[source]# Elementwise square: \(x^2\). Parameters: x (Union[Array, ndarray, bool_, number, bool, int, float, complex]) Return type: Array
