burn_tensor/tensor/api/
cast.rs

1use burn_backend::ops::{BoolTensorOps, FloatTensorOps, IntTensorOps};
2use burn_backend::{DType, FloatDType, IntDType};
3use burn_dispatch::Dispatch;
4
5use crate::kind::Basic;
6use crate::ops::BridgeTensor;
7use crate::{Bool, Float, Int, Tensor};
8
9/// Trait for types that represent a valid cast target from a tensor of kind `K`.
10///
11/// The generic parameter `K` is the *input* tensor kind ([`Float`], [`Int`], or [`Bool`]).
12/// Implementors declare the output kind and provide the actual cast logic.
13pub trait Cast<const D: usize, K: Basic> {
14    /// The output tensor kind after casting.
15    type OutputKind: Basic;
16
17    /// Cast a tensor primitive to the target dtype.
18    fn cast(tensor: Tensor<D, K>, dtype: Self) -> Tensor<D, Self::OutputKind>;
19}
20
21// --- Float input impls ---
22
23impl<const D: usize> Cast<D, Float> for FloatDType {
24    type OutputKind = Float;
25
26    fn cast(tensor: Tensor<D, Float>, dtype: Self) -> Tensor<D, Float> {
27        if tensor.primitive.is_float() {
28            let current: FloatDType = tensor.dtype().into();
29            if current == dtype {
30                return tensor;
31            }
32            Tensor::new(float_cast_impl(tensor.primitive, dtype))
33        } else {
34            panic!("Should be Float primitive kind");
35        }
36    }
37}
38
39impl<const D: usize> Cast<D, Float> for IntDType {
40    type OutputKind = Int;
41
42    fn cast(tensor: Tensor<D, Float>, dtype: Self) -> Tensor<D, Int> {
43        Tensor::new(float_to_int_impl(tensor.primitive, dtype))
44    }
45}
46
47/// Backward-compatible impl: only float `DType` variants are accepted.
48///
49/// # Panics
50///
51/// Panics if `dtype` is not a float variant (e.g., `DType::I32`).
52/// Use [`IntDType`] directly for cross-kind casting to int.
53impl<const D: usize> Cast<D, Float> for DType {
54    type OutputKind = Float;
55
56    fn cast(tensor: Tensor<D, Float>, dtype: Self) -> Tensor<D, Float> {
57        let float_dtype: FloatDType = dtype.into();
58        <FloatDType as Cast<D, Float>>::cast(tensor, float_dtype)
59    }
60}
61
62// --- Int input impls ---
63
64impl<const D: usize> Cast<D, Int> for IntDType {
65    type OutputKind = Int;
66
67    fn cast(tensor: Tensor<D, Int>, dtype: Self) -> Tensor<D, Int> {
68        let current: IntDType = tensor.primitive.dtype().into();
69        if current == dtype {
70            return tensor;
71        }
72        Tensor::new(int_cast_impl(tensor.primitive, dtype))
73    }
74}
75
76impl<const D: usize> Cast<D, Int> for FloatDType {
77    type OutputKind = Float;
78
79    fn cast(tensor: Tensor<D, Int>, dtype: Self) -> Tensor<D, Float> {
80        Tensor::new(int_to_float_impl(tensor.primitive, dtype))
81    }
82}
83
84/// Backward-compatible impl: only int `DType` variants are accepted.
85///
86/// # Panics
87///
88/// Panics if `dtype` is not an int variant (e.g., `DType::F32`).
89/// Use [`FloatDType`] directly for cross-kind casting to float.
90impl<const D: usize> Cast<D, Int> for DType {
91    type OutputKind = Int;
92
93    fn cast(tensor: Tensor<D, Int>, dtype: Self) -> Tensor<D, Int> {
94        let int_dtype: IntDType = dtype.into();
95        <IntDType as Cast<D, Int>>::cast(tensor, int_dtype)
96    }
97}
98
99// --- Bool input impls ---
100
101impl<const D: usize> Cast<D, Bool> for IntDType {
102    type OutputKind = Int;
103
104    fn cast(tensor: Tensor<D, Bool>, dtype: Self) -> Tensor<D, Int> {
105        Tensor::new(bool_cast_to_int_impl(tensor.primitive, dtype))
106    }
107}
108
109impl<const D: usize> Cast<D, Bool> for FloatDType {
110    type OutputKind = Float;
111
112    fn cast(tensor: Tensor<D, Bool>, dtype: Self) -> Tensor<D, Float> {
113        Tensor::new(bool_cast_to_float_impl(tensor.primitive, dtype))
114    }
115}
116
117// =========================================================================
118// Non-generic implementation helpers (outlined from the generic API).
119// See the crate-level docs for the rationale behind this pattern.
120// =========================================================================
121
122fn float_cast_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor {
123    BridgeTensor::float(Dispatch::float_cast(p.into_float(), dtype))
124}
125
126fn float_to_int_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor {
127    BridgeTensor::int(Dispatch::float_into_int(p.into_float(), dtype))
128}
129
130fn int_cast_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor {
131    BridgeTensor::int(Dispatch::int_cast(p.into(), dtype))
132}
133
134fn int_to_float_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor {
135    BridgeTensor::float(Dispatch::int_into_float(p.into(), dtype))
136}
137
138fn bool_cast_to_int_impl(p: BridgeTensor, dtype: IntDType) -> BridgeTensor {
139    BridgeTensor::bool(Dispatch::bool_into_int(p.into(), dtype))
140}
141
142fn bool_cast_to_float_impl(p: BridgeTensor, dtype: FloatDType) -> BridgeTensor {
143    BridgeTensor::float(Dispatch::bool_into_float(p.into(), dtype))
144}
burn_tensor/tensor/api/cast.rs

burn_tensor/tensor/api/
cast.rs
