jax.lax.square

jax.lax.square(x)[source]

Elementwise square: \(x^2\).

Parameters

x (Any) –

Return type

Any