Skip to contents

Element-wise selection based on a boolean predicate, like R's ifelse(). For each element, returns the corresponding element from true_value where pred is TRUE and from false_value where pred is FALSE.

Usage

prim_ifelse(pred, true_value, false_value)

Arguments

pred

(arrayish of boolean type)
Predicate array. Must be scalar or have the same shape as true_value.

true_value, false_value

(arrayish)
Values to select from. Must have the same dtype and shape.

Value

arrayish
Has the same dtype and shape as true_value. It is ambiguous if both true_value and false_value are ambiguous.

Implemented Rules

  • stablehlo

  • quickr

  • reverse

StableHLO

Lowers to stablehlo::hlo_select().

See also

Examples

pred <- nv_array(c(TRUE, FALSE, TRUE))
prim_ifelse(pred, nv_array(c(1, 2, 3)), nv_array(c(4, 5, 6)))
#> AnvlArray
#>  1
#>  5
#>  3
#> [ CPUf32{3} ]