Skip to contents

Selects elements from true_value or false_value based on pred, analogous to R's ifelse().

Usage

nv_ifelse(pred, true_value, false_value)

Arguments

pred

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

true_value

(tensorish)
Values to return where pred is TRUE.

false_value

(tensorish)
Values to return where pred is FALSE. Must have the same shape and data type as true_value.

Value

tensorish
Has the same shape and data type as true_value.

See also

nvl_ifelse() for the underlying primitive.

Examples

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