jax.lax.select

Warning

This page was created from a pull request (#9655).

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

previous

jax.lax.scatter_mul

next

jax.lax.shift_left