luma_tensor/ops/
cast.rs

1//! Type conversion between kinds and precisions.
2//!
3//! `Float::cast` (precision change within the float kind) records `Op::Cast` so
4//! its gradient flows back — the capability the old typed design could not express.
5
6use crate::{
7    Bool, DTypeKind, Device, Float, FloatMeta, Int, Tensor,
8    dtype::{BoolDType, FloatDType, IntDType},
9};
10
11pub trait Cast<D: Device, K: DTypeKind<D>> {
12    type Output: DTypeKind<D>;
13    fn cast(tensor: &Tensor<D, K>, dtype: Self) -> crate::Result<Tensor<D, Self::Output>>;
14}
15
16impl<D: Device, K: DTypeKind<D>> Tensor<D, K> {
17    pub fn cast<T: Cast<D, K>>(&self, dtype: T) -> crate::Result<Tensor<D, T::Output>> {
18        T::cast(self, dtype)
19    }
20}
21
22impl<D: Device, K: CastDTypeKind<D>> Cast<D, K> for FloatDType {
23    type Output = Float;
24    fn cast(tensor: &Tensor<D, K>, dtype: Self) -> crate::Result<Tensor<D, Self::Output>> {
25        tensor.cast_float(dtype)
26    }
27}
28
29impl<D: Device, K: CastDTypeKind<D>> Cast<D, K> for IntDType {
30    type Output = Int;
31    fn cast(tensor: &Tensor<D, K>, dtype: Self) -> crate::Result<Tensor<D, Self::Output>> {
32        tensor.cast_int(dtype)
33    }
34}
35
36impl<D: Device, K: CastDTypeKind<D>> Cast<D, K> for BoolDType {
37    type Output = Bool;
38    fn cast(tensor: &Tensor<D, K>, dtype: Self) -> crate::Result<Tensor<D, Self::Output>> {
39        tensor.cast_bool(dtype)
40    }
41}
42
43pub trait CastDTypeKind<D: Device>: DTypeKind<D> {
44    fn cast_float(tensor: &Tensor<D, Self>, to: FloatDType) -> crate::Result<Tensor<D, Float>>;
45    fn cast_int(tensor: &Tensor<D, Self>, to: IntDType) -> crate::Result<Tensor<D, Int>>;
46    fn cast_bool(tensor: &Tensor<D, Self>, to: BoolDType) -> crate::Result<Tensor<D, Bool>>;
47}
48
49impl<D: Device> CastDTypeKind<D> for Float {
50    fn cast_float(tensor: &Tensor<D, Self>, to: FloatDType) -> crate::Result<Tensor<D, Float>> {
51        let storage = D::f_cast_float(&*tensor.storage_read()?, tensor.layout(), to)?;
52        let meta = FloatMeta::on_cast(tensor);
53        Ok(Tensor::from_storage(storage, tensor.shape().clone(), meta))
54    }
55
56    fn cast_int(tensor: &Tensor<D, Self>, to: IntDType) -> crate::Result<Tensor<D, Int>> {
57        let storage = D::f_cast_int(&*tensor.storage_read()?, tensor.layout(), to)?;
58        Ok(Tensor::from_storage(storage, tensor.shape().clone(), ()))
59    }
60
61    fn cast_bool(tensor: &Tensor<D, Self>, to: BoolDType) -> crate::Result<Tensor<D, Bool>> {
62        let storage = D::f_cast_bool(&*tensor.storage_read()?, tensor.layout(), to)?;
63        Ok(Tensor::from_storage(storage, tensor.shape().clone(), ()))
64    }
65}
66
67impl<D: Device> CastDTypeKind<D> for Int {
68    fn cast_float(tensor: &Tensor<D, Self>, to: FloatDType) -> crate::Result<Tensor<D, Float>> {
69        let storage = D::i_cast_float(&*tensor.storage_read()?, tensor.layout(), to)?;
70        Ok(Tensor::from_storage(storage, tensor.shape().clone(), FloatMeta::val()))
71    }
72
73    fn cast_int(tensor: &Tensor<D, Self>, to: IntDType) -> crate::Result<Tensor<D, Int>> {
74        let storage = D::i_cast_int(&*tensor.storage_read()?, tensor.layout(), to)?;
75        Ok(Tensor::from_storage(storage, tensor.shape().clone(), ()))
76    }
77
78    fn cast_bool(tensor: &Tensor<D, Self>, to: BoolDType) -> crate::Result<Tensor<D, Bool>> {
79        let storage = D::i_cast_bool(&*tensor.storage_read()?, tensor.layout(), to)?;
80        Ok(Tensor::from_storage(storage, tensor.shape().clone(), ()))
81    }
82}
83
84impl<D: Device> CastDTypeKind<D> for Bool {
85    fn cast_float(tensor: &Tensor<D, Self>, to: FloatDType) -> crate::Result<Tensor<D, Float>> {
86        let storage = D::b_cast_float(&*tensor.storage_read()?, tensor.layout(), to)?;
87        Ok(Tensor::from_storage(storage, tensor.shape().clone(), FloatMeta::val()))
88    }
89
90    fn cast_int(tensor: &Tensor<D, Self>, to: IntDType) -> crate::Result<Tensor<D, Int>> {
91        let storage = D::b_cast_int(&*tensor.storage_read()?, tensor.layout(), to)?;
92        Ok(Tensor::from_storage(storage, tensor.shape().clone(), ()))
93    }
94
95    fn cast_bool(tensor: &Tensor<D, Self>, to: BoolDType) -> crate::Result<Tensor<D, Bool>> {
96        let storage = D::b_cast_bool(&*tensor.storage_read()?, tensor.layout(), to)?;
97        Ok(Tensor::from_storage(storage, tensor.shape().clone(), ()))
98    }
99}
100
101impl<D: Device, K: CastDTypeKind<D>> Tensor<D, K> {
102    #[inline]
103    pub fn cast_float(&self, to: FloatDType) -> crate::Result<Tensor<D, Float>> {
104        K::cast_float(self, to)
105    }
106
107    #[inline]
108    pub fn cast_int(&self, to: IntDType) -> crate::Result<Tensor<D, Int>> {
109        K::cast_int(self, to)
110    }
111
112    #[inline]
113    pub fn cast_bool(&self, to: BoolDType) -> crate::Result<Tensor<D, Bool>> {
114        K::cast_bool(self, to)
115    }
116}
luma_tensor/ops/cast.rs

luma_tensor/ops/
cast.rs
