jax.lax.select#

jax.lax.select(pred, on_true, on_false)[source]#

Wraps XLA’s Select operator.

Parameters

  • pred (Any) –

  • on_true (Any) –

  • on_false (Any) –

Return type

Any