burn-tensor 0.22.0-pre.1

Tensor library with user-friendly APIs and automatic differentiation support
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144

use burn_backend::ops::{BoolTensorOps, FloatTensorOps, IntTensorOps};
use burn_backend::{DType, FloatDType, IntDType};
use burn_dispatch::Dispatch;

use crate::kind::Basic;
use crate::ops::BridgeTensor;
use crate::{Bool, Float, Int, Tensor};

/// Trait for types that represent a valid cast target from a tensor of kind `K`.
///
/// The generic parameter `K` is the *input* tensor kind ([`Float`], [`Int`], or [`Bool`]).
/// Implementors declare the output kind and provide the actual cast logic.
pub trait Cast<const D: usize, K: Basic> {
    /// The output tensor kind after casting.
    type OutputKind: Basic;

    /// Cast a tensor primitive to the target dtype.
    fn cast(tensor: Tensor<D, K>, dtype: Self) -> Tensor<D, Self::OutputKind>;
}

// --- Float input impls ---

impl<const D: usize> Cast<D, Float> for FloatDType {
    type OutputKind = Float;

    fn cast(tensor: Tensor<D, Float>, dtype: Self) -> Tensor<D, Float> {
        if tensor.primitive.is_float() {
            let current: FloatDType = tensor.dtype().into();
            if current == dtype {
                return tensor;
            }
            Tensor::new(float_cast_impl(tensor.primitive, dtype))
        } else {
            panic!("Should be Float primitive kind");
        }
    }
}

impl<const D: usize> Cast<D, Float> for IntDType {
    type OutputKind = Int;

    fn cast(tensor: Tensor<D, Float>, dtype: Self) -> Tensor<D, Int> {
        Tensor::new(float_to_int_impl(tensor.primitive, dtype))
    }
}

/// Backward-compatible impl: only float `DType` variants are accepted.
///
/// # Panics
///
/// Panics if `dtype` is not a float variant (e.g., `DType::I32`).
/// Use [`IntDType`] directly for cross-kind casting to int.
impl<const D: usize> Cast<D, Float> for DType {
    type OutputKind = Float;

    fn cast(tensor: Tensor<D, Float>, dtype: Self) -> Tensor<D, Float> {
        let float_dtype: FloatDType = dtype.into();
        <FloatDType as Cast<D, Float>>::cast(tensor, float_dtype)
    }
}

// --- Int input impls ---

impl<const D: usize> Cast<D, Int> for IntDType {
    type OutputKind = Int;

    fn cast(tensor: Tensor<D, Int>, dtype: Self) -> Tensor<D, Int> {
        let current: IntDType = tensor.primitive.dtype().into();
        if current == dtype {
            return tensor;
        }
        Tensor::new(int_cast_impl(tensor.primitive, dtype))
    }
}

impl<const D: usize> Cast<D, Int> for FloatDType {
    type OutputKind = Float;

    fn cast(tensor: Tensor<D, Int>, dtype: Self) -> Tensor<D, Float> {
        Tensor::new(int_to_float_impl(tensor.primitive, dtype))
    }
}

/// Backward-compatible impl: only int `DType` variants are accepted.
///
/// # Panics
///
/// Panics if `dtype` is not an int variant (e.g., `DType::F32`).
/// Use [`FloatDType`] directly for cross-kind casting to float.
impl<const D: usize> Cast<D, Int> for DType {
    type OutputKind = Int;

    fn cast(tensor: Tensor<D, Int>, dtype: Self) -> Tensor<D, Int> {
        let int_dtype: IntDType = dtype.into();
        <IntDType as Cast<D, Int>>::cast(tensor, int_dtype)
    }
}

// --- Bool input impls ---

impl<const D: usize> Cast<D, Bool> for IntDType {
    type OutputKind = Int;

    fn cast(tensor: Tensor<D, Bool>, dtype: Self) -> Tensor<D, Int> {
        Tensor::new(bool_cast_to_int_impl(tensor.primitive, dtype))
    }
}

impl<const D: usize> Cast<D, Bool> for FloatDType {
    type OutputKind = Float;

    fn cast(tensor: Tensor<D, Bool>, dtype: Self) -> Tensor<D, Float> {
        Tensor::new(bool_cast_to_float_impl(tensor.primitive, dtype))
    }
}

// =========================================================================
// Non-generic implementation helpers (outlined from the generic API).
// See the crate-level docs for the rationale behind this pattern.
// =========================================================================

fn float_cast_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor {
    BridgeTensor::float(Dispatch::float_cast(p.into_float(), dtype))
}

fn float_to_int_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor {
    BridgeTensor::int(Dispatch::float_into_int(p.into_float(), dtype))
}

fn int_cast_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor {
    BridgeTensor::int(Dispatch::int_cast(p.into(), dtype))
}

fn int_to_float_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor {
    BridgeTensor::float(Dispatch::int_into_float(p.into(), dtype))
}

fn bool_cast_to_int_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor {
    BridgeTensor::bool(Dispatch::bool_into_int(p.into(), dtype))
}

fn bool_cast_to_float_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor {
    BridgeTensor::float(Dispatch::bool_into_float(p.into(), dtype))
}