jax.lax.square¶

jax.lax.square(x)[source]¶

Elementwise square: \(x^2\).

Parameters

x (Any) –

Return type

Any