sim_lib_numbers_tensor/implementation/
cast.rs

1//! Explicit tensor dtype conversion with checked narrowing and deterministic
2//! rounding.
3
4use std::sync::Arc;
5
6use half::{bf16, f16};
7use sim_kernel::{Cx, Error, Expr, QuoteMode, Result, Symbol, Value};
8
9use crate::{
10    Tensor, TypedTensorStorage, domains, number_literal_for_tensor_cell, tensor_value_ref,
11};
12
13/// Returns the qualified function symbol for explicit tensor dtype conversion.
14pub fn cast_symbol() -> Symbol {
15    Symbol::qualified("tensor", "cast")
16}
17
18/// Casts a tensor into `target_dtype`.
19///
20/// Supported targets are `numbers/i64`, `numbers/f32`, `numbers/f64`,
21/// `numbers/f16`, and `numbers/bf16`. Floating-to-integer casts round ties to
22/// even and reject NaN, infinity, and out-of-range results. Floating narrowing
23/// preserves NaN, infinity, and signed zero, and rejects finite overflow into an
24/// infinite result.
25pub fn cast_tensor(tensor: &Tensor, target_dtype: Symbol) -> Result<Tensor> {
26    if tensor.dtype() == &target_dtype {
27        return Ok(tensor.clone());
28    }
29    let cells = tensor
30        .cells()?
31        .iter()
32        .enumerate()
33        .map(|(index, value)| CastCell::from_value(value, index))
34        .collect::<Result<Vec<_>>>()?;
35    let shape = tensor.shape().to_vec();
36    if target_dtype == domains::i64() {
37        let out = cells
38            .iter()
39            .enumerate()
40            .map(|(index, cell)| cell.to_i64(index))
41            .collect::<Result<Vec<_>>>()?;
42        return Tensor::from_storage(
43            shape,
44            target_dtype,
45            Arc::new(TypedTensorStorage::<i64>::new(out)),
46        );
47    }
48    if target_dtype == domains::f32() {
49        let out = cells
50            .iter()
51            .enumerate()
52            .map(|(index, cell)| cell.to_f32(index, domains::f32()))
53            .collect::<Result<Vec<_>>>()?;
54        return Tensor::from_storage(
55            shape,
56            target_dtype,
57            Arc::new(TypedTensorStorage::<f32>::new(out)),
58        );
59    }
60    if target_dtype == domains::f64() {
61        let out = cells
62            .iter()
63            .copied()
64            .map(CastCell::to_f64)
65            .collect::<Vec<_>>();
66        return Tensor::from_storage(
67            shape,
68            target_dtype,
69            Arc::new(TypedTensorStorage::<f64>::new(out)),
70        );
71    }
72    if target_dtype == domains::f16() {
73        let out = cells
74            .iter()
75            .enumerate()
76            .map(|(index, cell)| cell.to_f16(index))
77            .collect::<Result<Vec<_>>>()?;
78        return Tensor::from_storage(
79            shape,
80            target_dtype,
81            Arc::new(TypedTensorStorage::<f16>::new(out)),
82        );
83    }
84    if target_dtype == domains::bf16() {
85        let out = cells
86            .iter()
87            .enumerate()
88            .map(|(index, cell)| cell.to_bf16(index))
89            .collect::<Result<Vec<_>>>()?;
90        return Tensor::from_storage(
91            shape,
92            target_dtype,
93            Arc::new(TypedTensorStorage::<bf16>::new(out)),
94        );
95    }
96    Err(Error::Eval(format!(
97        "tensor/cast does not support target dtype {target_dtype}"
98    )))
99}
100
101/// Casts a tensor value into `target_dtype` and returns it as a runtime value.
102pub fn cast_tensor_value(cx: &mut Cx, value: Value, target_dtype: Symbol) -> Result<Value> {
103    let tensor = tensor_value_ref(&value)
104        .ok_or_else(|| Error::Eval("tensor/cast expects a tensor value".to_owned()))?;
105    if tensor.dtype() == &target_dtype {
106        return Ok(value);
107    }
108    cx.factory()
109        .opaque(Arc::new(cast_tensor(tensor, target_dtype)?))
110}
111
112pub(crate) fn cast_function_impl(cx: &mut Cx, values: Vec<Value>) -> Result<Value> {
113    let [tensor_value, dtype_value] = values.as_slice() else {
114        return Err(Error::Eval(
115            "tensor/cast expects exactly two arguments: tensor dtype".to_owned(),
116        ));
117    };
118    let dtype = extract_dtype_symbol(cx, dtype_value)?;
119    cast_tensor_value(cx, tensor_value.clone(), dtype)
120}
121
122#[derive(Clone, Copy)]
123enum CastCell {
124    I64(i64),
125    F32(f32),
126    F64(f64),
127    F16(f16),
128    Bf16(bf16),
129}
130
131impl CastCell {
132    fn from_value(value: &Value, index: usize) -> Result<Self> {
133        let literal = number_literal_for_tensor_cell(value).ok_or_else(|| {
134            Error::Eval(format!(
135                "tensor/cast source cell {index} does not expose a numeric literal"
136            ))
137        })?;
138        if literal.domain == domains::i64() {
139            return Ok(Self::I64(parse_literal(
140                &literal.canonical,
141                index,
142                &literal.domain,
143            )?));
144        }
145        if literal.domain == domains::f32() {
146            return Ok(Self::F32(parse_literal(
147                &literal.canonical,
148                index,
149                &literal.domain,
150            )?));
151        }
152        if literal.domain == domains::f64() {
153            return Ok(Self::F64(parse_literal(
154                &literal.canonical,
155                index,
156                &literal.domain,
157            )?));
158        }
159        if literal.domain == domains::f16() {
160            let value = parse_literal::<f32>(&literal.canonical, index, &literal.domain)?;
161            return Ok(Self::F16(f16::from_f32(value)));
162        }
163        if literal.domain == domains::bf16() {
164            let value = parse_literal::<f32>(&literal.canonical, index, &literal.domain)?;
165            return Ok(Self::Bf16(bf16::from_f32(value)));
166        }
167        Err(Error::Eval(format!(
168            "tensor/cast does not support source dtype {} at cell {index}",
169            literal.domain
170        )))
171    }
172
173    fn to_f64(self) -> f64 {
174        match self {
175            Self::I64(value) => value as f64,
176            Self::F32(value) => f64::from(value),
177            Self::F64(value) => value,
178            Self::F16(value) => f64::from(value.to_f32()),
179            Self::Bf16(value) => f64::from(value.to_f32()),
180        }
181    }
182
183    fn to_f32(self, index: usize, target: Symbol) -> Result<f32> {
184        match self {
185            Self::I64(value) => Ok(value as f32),
186            Self::F32(value) => Ok(value),
187            Self::F64(value) => {
188                let narrowed = value as f32;
189                reject_finite_overflow(value, narrowed.is_infinite(), index, &target)?;
190                Ok(narrowed)
191            }
192            Self::F16(value) => Ok(value.to_f32()),
193            Self::Bf16(value) => Ok(value.to_f32()),
194        }
195    }
196
197    fn to_i64(self, index: usize) -> Result<i64> {
198        let value = self.to_f64();
199        if !value.is_finite() {
200            return Err(Error::Eval(format!(
201                "tensor/cast cannot cast non-finite cell {index} to {}",
202                domains::i64()
203            )));
204        }
205        let rounded = value.round_ties_even();
206        const I64_MIN_INCLUSIVE: f64 = -9_223_372_036_854_775_808.0;
207        const I64_MAX_EXCLUSIVE: f64 = 9_223_372_036_854_775_808.0;
208        if !(I64_MIN_INCLUSIVE..I64_MAX_EXCLUSIVE).contains(&rounded) {
209            return Err(Error::Eval(format!(
210                "tensor/cast cell {index} overflows {}",
211                domains::i64()
212            )));
213        }
214        Ok(rounded as i64)
215    }
216
217    fn to_f16(self, index: usize) -> Result<f16> {
218        let source = self.to_f64();
219        let narrowed = f16::from_f32(self.to_f32(index, domains::f16())?);
220        reject_finite_overflow(source, narrowed.is_infinite(), index, &domains::f16())?;
221        Ok(narrowed)
222    }
223
224    fn to_bf16(self, index: usize) -> Result<bf16> {
225        let source = self.to_f64();
226        let narrowed = bf16::from_f32(self.to_f32(index, domains::bf16())?);
227        reject_finite_overflow(source, narrowed.is_infinite(), index, &domains::bf16())?;
228        Ok(narrowed)
229    }
230}
231
232fn reject_finite_overflow(
233    source: f64,
234    narrowed_is_infinite: bool,
235    index: usize,
236    target: &Symbol,
237) -> Result<()> {
238    if source.is_finite() && narrowed_is_infinite {
239        return Err(Error::Eval(format!(
240            "tensor/cast cell {index} overflows {target}"
241        )));
242    }
243    Ok(())
244}
245
246fn parse_literal<T: std::str::FromStr>(canonical: &str, index: usize, domain: &Symbol) -> Result<T>
247where
248    T::Err: std::fmt::Display,
249{
250    canonical.parse::<T>().map_err(|err| {
251        Error::Eval(format!(
252            "tensor/cast cell {index} in {domain} has invalid canonical literal {canonical:?}: {err}"
253        ))
254    })
255}
256
257fn extract_dtype_symbol(cx: &mut Cx, value: &Value) -> Result<Symbol> {
258    match value.object().as_expr(cx)? {
259        Expr::Symbol(symbol) => Ok(symbol),
260        Expr::Quote {
261            mode: QuoteMode::Quote,
262            expr,
263        } => match *expr {
264            Expr::Symbol(symbol) => Ok(symbol),
265            _ => Err(Error::Eval(
266                "tensor/cast expected a symbol dtype".to_owned(),
267            )),
268        },
269        _ => Err(Error::Eval(
270            "tensor/cast expected a symbol dtype".to_owned(),
271        )),
272    }
273}
sim_lib_numbers_tensor/implementation/cast.rs

sim_lib_numbers_tensor/implementation/
cast.rs
