burn-jit 0.16.1

Generic backend that can be compiled just-in-time to any shader language target
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34

use super::{mask_where::MaskWhereStrategy, MaskFillStrategy};
use crate::{element::JitElement, tensor::JitTensor, BoolElement, JitRuntime};

/// Execute the mask fill kernel.
pub(crate) fn mask_fill_auto<R: JitRuntime, E: JitElement, BT: BoolElement>(
    tensor: JitTensor<R>,
    mask: JitTensor<R>,
    value: E,
) -> JitTensor<R> {
    let strategy = if tensor.can_mut() {
        MaskFillStrategy::Inplace
    } else {
        MaskFillStrategy::Readonly
    };

    super::mask_fill::<R, E, BT>(tensor, mask, value, strategy)
}

/// Execute the mask where kernel.
pub(crate) fn mask_where_auto<R: JitRuntime, E: JitElement, BT: BoolElement>(
    tensor: JitTensor<R>,
    mask: JitTensor<R>,
    value: JitTensor<R>,
) -> JitTensor<R> {
    let strategy = if tensor.can_mut_broadcast(&value) {
        MaskWhereStrategy::InplaceLhs
    } else if value.can_mut_broadcast(&tensor) {
        MaskWhereStrategy::InplaceRhs
    } else {
        MaskWhereStrategy::Readonly
    };

    super::mask_where::<R, E, BT>(tensor, mask, value, strategy)
}