Skip to contents

Primitive Ifelse

Source: R/primitives.R
nvl_ifelse.Rd

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

nvl_ifelse(pred, true_value, false_value)

Arguments

pred

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

true_value, false_value

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

Value

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

Implemented Rules

  • stablehlo

  • backward

StableHLO

Lowers to stablehlo::hlo_select().

See also

nv_ifelse()

Examples

jit_eval({
  pred <- nv_tensor(c(TRUE, FALSE, TRUE))
  nvl_ifelse(pred, nv_tensor(c(1, 2, 3)), nv_tensor(c(4, 5, 6)))
})
#> AnvilTensor
#>  1
#>  5
#>  3
#> [ CPUf32{3} ]