burn-candle 0.20.1

Candle backend for the Burn framework
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
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

use burn_backend::{DType, FloatDType, IntDType, Shape, quantization::QuantScheme};
use burn_backend::{Element, QTensorPrimitive, TensorData, TensorMetadata};

use crate::{CandleDevice, element::CandleElement};

/// A tensor that uses the candle backend.
#[derive(Debug, Clone)]
pub struct CandleTensor {
    pub(crate) tensor: candle_core::Tensor,
}

impl TensorMetadata for CandleTensor {
    fn dtype(&self) -> DType {
        match self.tensor.dtype() {
            candle_core::DType::U8 => DType::U8,
            candle_core::DType::U32 => DType::U32,
            candle_core::DType::I64 => DType::I64,
            candle_core::DType::BF16 => DType::BF16,
            candle_core::DType::F16 => DType::F16,
            candle_core::DType::F32 => DType::F32,
            candle_core::DType::F64 => DType::F64,
        }
    }

    fn shape(&self) -> Shape {
        Shape::from(self.tensor.dims().to_vec())
    }

    fn rank(&self) -> usize {
        self.tensor.dims().len()
    }
}

impl QTensorPrimitive for CandleTensor {
    fn scheme(&self) -> &QuantScheme {
        unimplemented!("Quantization is not supported")
    }
}

impl CandleTensor {
    /// Create a new tensor.
    pub fn new(tensor: candle_core::Tensor) -> Self {
        Self { tensor }
    }

    /// Creates a new tensor from data and a device.
    ///
    /// # Arguments
    ///
    /// * `data` - The tensor's data.
    /// * `device` - The device on which the tensor will be allocated.
    ///
    /// # Returns
    ///
    /// A new tensor.
    pub fn from_data<E: CandleElement>(data: TensorData, device: CandleDevice) -> Self {
        let candle_shape: candle_core::Shape = data.shape.clone().into();
        let tensor = candle_core::Tensor::from_slice(
            data.as_slice::<E>().unwrap(),
            candle_shape,
            &device.into(),
        );
        Self::new(tensor.unwrap())
    }
}

pub(crate) trait IntoDType {
    fn try_into_dtype(self) -> Result<candle_core::DType, candle_core::Error>;

    fn into_dtype(self) -> candle_core::DType
    where
        Self: Sized,
    {
        self.try_into_dtype().unwrap()
    }
}

impl IntoDType for IntDType {
    fn try_into_dtype(self) -> Result<candle_core::DType, candle_core::Error> {
        let dtype: DType = self.into();
        dtype.try_into_dtype()
    }
}

impl IntoDType for FloatDType {
    fn try_into_dtype(self) -> Result<candle_core::DType, candle_core::Error> {
        let dtype: DType = self.into();
        dtype.try_into_dtype()
    }
}

impl IntoDType for DType {
    fn try_into_dtype(self) -> Result<candle_core::DType, candle_core::Error> {
        match self {
            DType::F64 => Ok(candle_core::DType::F64),
            DType::F32 => Ok(candle_core::DType::F32),
            DType::Flex32 => Ok(candle_core::DType::F32),
            DType::F16 => Ok(candle_core::DType::F16),
            DType::BF16 => Ok(candle_core::DType::BF16),
            DType::I64 => Ok(candle_core::DType::I64),
            DType::U32 => Ok(candle_core::DType::U32),
            DType::U8 => Ok(candle_core::DType::U8),
            // DType::Bool => Ok(candle_core::DType::U8),
            _ => Err(candle_core::Error::Msg(format!(
                "Unsupported dtype {self:?}"
            ))),
        }
    }
}