maidenx_tensor 0.1.5

maidenx tensor
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

use crate::Tensor;
#[cfg(feature = "mps")]
use maidenx_core::device::Device;
use maidenx_core::{dtype::DType, error::Result, scalar::Scalar};

pub fn get_signed_dtype(dtype: DType) -> DType {
    match dtype {
        DType::U8 => DType::I32,
        DType::U16 => DType::I32,
        DType::U32 => DType::I64,
        DType::U64 => DType::I64,
        _ => dtype,
    }
}

pub fn get_promoted_dtype(dtype1: DType, dtype2: DType) -> DType {
    let mut dtype1 = dtype1;
    let mut dtype2 = dtype2;

    if dtype1 == DType::BOOL {
        dtype1 = DType::U8;
    }
    if dtype2 == DType::BOOL {
        dtype2 = DType::U8;
    }

    match (dtype1, dtype2) {
        (dtype1, dtype2) if dtype1 == dtype2 => dtype1,

        (_, DType::F64) | (DType::F64, _) => DType::F64,
        (_, DType::F32) | (DType::F32, _) => DType::F32,
        (DType::BF16, DType::F16) | (DType::F16, DType::BF16) => DType::F32,
        (_, DType::F16) | (DType::F16, _) => DType::F16,
        (_, DType::BF16) | (DType::BF16, _) => DType::BF16,

        (_, DType::I64) | (DType::I64, _) => DType::I64,
        (_, DType::I32) | (DType::I32, _) => DType::I32,
        (_, DType::I16) | (DType::I16, _) => DType::I16,
        (_, DType::I8) | (DType::I8, _) => DType::I8,
        (_, DType::U64) | (DType::U64, _) => DType::I64,
        (_, DType::U32) | (DType::U32, _) => DType::I64,
        (_, DType::U16) | (DType::U16, _) => DType::I32,
        (_, DType::U8) | (DType::U8, _) => DType::I32,

        _ => dtype1,
    }
}

// pub fn get_promoted_dtype_with_scalar(tensor_dtype: DType, scalar_dtype: DType) -> DType {
//     let tensor_dtype = if tensor_dtype == DType::BOOL { DType::U8 } else { tensor_dtype };
//
//     let scalar_dtype = if scalar_dtype == DType::BOOL { DType::U8 } else { scalar_dtype };
//
//     match (tensor_dtype, scalar_dtype) {
//         (tensor_dtype, scalar_dtype) if tensor_dtype == scalar_dtype => tensor_dtype,
//
//         (t_dtype, s_dtype) if t_dtype.is_int() && s_dtype.is_float() => s_dtype,
//         _ => tensor_dtype,
//     }
// }

pub fn promote_tensor(src: &Tensor, target_dtype: DType) -> Result<Tensor> {
    let src = if src.dtype() != target_dtype {
        let mut src = src.clone();
        src.with_dtype(target_dtype)?;

        src
    } else {
        src.clone()
    };

    Ok(src)
}

pub fn promote_scalar_for_tensor(src: Scalar, target_dtype: DType, with_tensor: &Tensor) -> Result<Scalar> {
    let src = if src.dtype() != target_dtype {
        let src = src.to_dtype(target_dtype);

        src
    } else {
        src.clone()
    };

    #[cfg(feature = "mps")]
    let src = if with_tensor.device() == Device::MPS {
        match src.dtype() {
            DType::U64 => src.to_dtype(DType::U32),
            DType::I64 => src.to_dtype(DType::I32),
            DType::F64 => src.to_dtype(DType::F32),
            _ => src,
        }
    } else {
        src
    };

    #[cfg(not(feature = "mps"))]
    let _ = with_tensor;

    Ok(src)
}