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
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
use alloc::vec;
use alloc::vec::Vec;
use burn_tensor::ops::{BoolTensorOps, IntTensorOps};
use core::ops::Range;
use crate::element::FloatNdArrayElement;
use crate::NdArrayDevice;
use crate::{tensor::NdArrayTensor, NdArrayBackend};
use burn_tensor::{backend::Backend, Data, Shape};
use super::NdArrayOps;
impl<E: FloatNdArrayElement> BoolTensorOps<NdArrayBackend<E>> for NdArrayBackend<E> {
fn bool_from_data<const D: usize>(
data: Data<bool, D>,
_device: &NdArrayDevice,
) -> NdArrayTensor<bool, D> {
NdArrayTensor::from_data(data)
}
fn bool_shape<const D: usize>(
tensor: &<NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
) -> Shape<D> {
tensor.shape()
}
fn bool_to_data<const D: usize>(
tensor: &<NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
) -> Data<bool, D> {
let values = tensor.array.iter().map(Clone::clone).collect();
Data::new(values, tensor.shape())
}
fn bool_into_data<const D: usize>(
tensor: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
) -> Data<bool, D> {
let shape = tensor.shape();
let values = tensor.array.into_iter().collect();
Data::new(values, shape)
}
fn bool_to_device<const D: usize>(
tensor: NdArrayTensor<bool, D>,
_device: &NdArrayDevice,
) -> NdArrayTensor<bool, D> {
tensor
}
fn bool_reshape<const D1: usize, const D2: usize>(
tensor: NdArrayTensor<bool, D1>,
shape: Shape<D2>,
) -> NdArrayTensor<bool, D2> {
NdArrayOps::reshape(tensor, shape)
}
fn bool_index<const D1: usize, const D2: usize>(
tensor: NdArrayTensor<bool, D1>,
indexes: [Range<usize>; D2],
) -> NdArrayTensor<bool, D1> {
NdArrayOps::index(tensor, indexes)
}
fn bool_into_int<const D: usize>(
tensor: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
) -> NdArrayTensor<i64, D> {
let data = Self::bool_into_data(tensor);
NdArrayBackend::<E>::int_from_data(data.convert(), &NdArrayDevice::Cpu)
}
fn bool_device<const D: usize>(
_tensor: &<NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
) -> <NdArrayBackend<E> as Backend>::Device {
NdArrayDevice::Cpu
}
fn bool_empty<const D: usize>(
shape: Shape<D>,
_device: &<NdArrayBackend<E> as Backend>::Device,
) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
let values = vec![false; shape.num_elements()];
NdArrayTensor::from_data(Data::new(values, shape))
}
fn bool_index_assign<const D1: usize, const D2: usize>(
tensor: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D1>,
indexes: [Range<usize>; D2],
value: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D1>,
) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D1> {
NdArrayOps::index_assign(tensor, indexes, value)
}
fn bool_cat<const D: usize>(
tensors: Vec<<NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>>,
dim: usize,
) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
NdArrayOps::cat(tensors, dim)
}
fn bool_equal<const D: usize>(
lhs: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
rhs: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
let mut array = lhs.array;
array.zip_mut_with(&rhs.array, |a, b| *a = *a && *b);
NdArrayTensor { array }
}
fn bool_equal_elem<const D: usize>(
lhs: <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D>,
rhs: bool,
) -> <NdArrayBackend<E> as Backend>::BoolTensorPrimitive<D> {
let array = lhs.array.mapv(|a| a == rhs).into_shared();
NdArrayTensor { array }
}
}