glowstick-burn 0.2.0

Integration of glowstick with the burn tensor
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
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71

use std::marker::PhantomData;

use burn::prelude::Backend;
use burn::tensor::{Bool, Tensor as BTensor};

use glowstick::Shape;

use crate::Tensor;

/// Creates a lower triangular mask of the specified shape.
///
/// # Example
///
/// ```rust
/// # fn main() -> Result<(), glowstick_burn::Error> {
/// # use burn::backend::ndarray::{NdArray, NdArrayDevice};
/// # type Backend = NdArray;
/// use burn::tensor::{Device, Tensor as BurnTensor};
/// use glowstick_burn::{tril_mask, Tensor};
/// use glowstick::{Shape2, num::*, dyndims};
///
/// dyndims! {
///     B: BatchSize,
///     N: SequenceLength
/// }
///
/// let device = NdArrayDevice::Cpu;
/// let mask = tril_mask!(0, &device, Backend, [{ 2 } => B, { 3 } => N]);
///
/// assert_eq!(
///     mask.inner().to_data().to_vec::<bool>().unwrap(),
///     &[
///         false, true, true,
///         false, false, true
///     ]
/// );
/// # Ok(())
/// # }
/// ```
#[macro_export]
macro_rules! tril_mask {
    ($o:expr,$d:expr,$b:ty,[$($ds:tt)+]) => {{
        use $crate::op::tril_mask::TrilMask;
        (
            std::marker::PhantomData::<$b>,
            std::marker::PhantomData::<glowstick::TensorShape<$crate::reshape_tys!($($ds)+)>>,
            $o,
            $d,
        )
            .tril_mask($crate::reshape_val!($($ds)+).into_array())
    }};
}

pub trait TrilMask<const M: usize> {
    type Out;
    fn tril_mask(self, shape: [usize; M]) -> Self::Out;
}
impl<B, S, const M: usize> TrilMask<M>
    for (PhantomData<B>, PhantomData<S>, i64, &<B as Backend>::Device)
where
    B: Backend,
    S: Shape,
{
    type Out = Tensor<BTensor<B, M, Bool>, S>;
    fn tril_mask(self, shape: [usize; M]) -> Self::Out {
        Tensor(
            BTensor::<B, M, Bool>::tril_mask(shape, self.2, self.3),
            PhantomData,
        )
    }
}