edgefirst_tensor/lens.rs
1// SPDX-FileCopyrightText: Copyright 2026 Au-Zone Technologies
2// SPDX-License-Identifier: Apache-2.0
3
4//! `Tensor<T>` as a typed lens over [`TensorDyn`].
5//!
6//! `TensorDyn` is the main type — the one every FFI surface (C, Python, and
7//! any future binding) calls `Tensor`. Rust code that needs the element type
8//! known at compile time opens a lens onto it with
9//! [`TensorDyn::as_typed`]/[`TensorDyn::as_typed_mut`] instead of matching a
10//! constructor; `Tensor<T>` itself owns no storage of its own to lose — it
11//! borrows the handle's.
12//!
13//! Rust-only. The lens exists because Rust can express "this handle's
14//! element type is `f32`" as a static fact; C and Python have no such
15//! concept and always operate on the erased `TensorDyn` directly. Both
16//! backends implement [`as_typed`](TensorDyn::as_typed)/
17//! [`as_typed_mut`](TensorDyn::as_typed_mut) -- `static_lens` below for the
18//! enum-of-`Tensor<T>` backend, `dynamic_lens` for the opaque-handle one --
19//! with different validity arguments for the same cast; see each module's
20//! docs.
21
22use crate::sealed::Sealed;
23use crate::DType;
24use half::f16;
25use num_traits::Num;
26use std::fmt;
27
28/// An element type a `Tensor<T>` lens may be opened for.
29///
30/// Sealed: the eleven variants of [`TensorDyn`] are the only element types a
31/// tensor can carry, and a downstream impl would let
32/// [`as_typed`](TensorDyn::as_typed) name a dtype no tensor can hold.
33pub trait Element: Sealed + Num + Clone + fmt::Debug + Send + Sync + 'static {
34 /// The runtime dtype this element type corresponds to.
35 const DTYPE: DType;
36}
37
38macro_rules! element {
39 ($($t:ty => $d:ident),* $(,)?) => {$(
40 impl Element for $t { const DTYPE: DType = DType::$d; }
41 )*};
42}
43element! {
44 u8 => U8, i8 => I8, u16 => U16, i16 => I16, u32 => U32, i32 => I32,
45 u64 => U64, i64 => I64, f16 => F16, f32 => F32, f64 => F64,
46}
47
48#[cfg(feature = "static")]
49mod static_lens {
50 //! The lens accessors, for the backend where `TensorDyn` really is an
51 //! enum over eleven `Tensor<T>` values.
52 //!
53 //! Implementation note: this deliberately does **not** reinterpret
54 //! `TensorDyn` as `Tensor<T>` via a raw pointer cast. `TensorDyn` is a
55 //! tagged-union `enum`, and Rust gives no layout guarantee that would
56 //! make `&TensorDyn as *const Tensor<T>` sound for any variant --
57 //! that's a promise only `#[repr(transparent)]` over a *single* field
58 //! makes, and `TensorDyn` has eleven. Instead this matches the active
59 //! variant (as `as_u8`/`as_f32`/etc. in `static_backend.rs` already do)
60 //! and uses `std::any::Any::downcast_ref`, which performs the exact
61 //! `TypeId` check that proves `T` is the concrete element type -- safe,
62 //! no `unsafe` block, and it costs nothing `as_u8()` wasn't already
63 //! paying.
64 use super::Element;
65 use crate::{Tensor, TensorDyn};
66 use std::any::Any;
67
68 impl TensorDyn {
69 /// Borrow this tensor as a statically-typed `Tensor<T>`, or `None`
70 /// when the element type does not match.
71 pub fn as_typed<T: Element>(&self) -> Option<&Tensor<T>> {
72 let any: &dyn Any = match self {
73 TensorDyn::U8(t) => t,
74 TensorDyn::I8(t) => t,
75 TensorDyn::U16(t) => t,
76 TensorDyn::I16(t) => t,
77 TensorDyn::U32(t) => t,
78 TensorDyn::I32(t) => t,
79 TensorDyn::U64(t) => t,
80 TensorDyn::I64(t) => t,
81 TensorDyn::F16(t) => t,
82 TensorDyn::F32(t) => t,
83 TensorDyn::F64(t) => t,
84 };
85 any.downcast_ref::<Tensor<T>>()
86 }
87
88 /// Mutably borrow this tensor as a statically-typed `Tensor<T>`, or
89 /// `None` when the element type does not match.
90 pub fn as_typed_mut<T: Element>(&mut self) -> Option<&mut Tensor<T>> {
91 let any: &mut dyn Any = match self {
92 TensorDyn::U8(t) => t,
93 TensorDyn::I8(t) => t,
94 TensorDyn::U16(t) => t,
95 TensorDyn::I16(t) => t,
96 TensorDyn::U32(t) => t,
97 TensorDyn::I32(t) => t,
98 TensorDyn::U64(t) => t,
99 TensorDyn::I64(t) => t,
100 TensorDyn::F16(t) => t,
101 TensorDyn::F32(t) => t,
102 TensorDyn::F64(t) => t,
103 };
104 any.downcast_mut::<Tensor<T>>()
105 }
106 }
107}
108
109#[cfg(feature = "dynamic")]
110mod dynamic_lens {
111 //! The lens accessors, for the backend where `TensorDyn` really is one
112 //! handle.
113 //!
114 //! Implementation note: unlike `static_lens`, this one *does*
115 //! reinterpret `&TensorDyn`/`&mut TensorDyn` as `&Tensor<T>`/`&mut
116 //! Tensor<T>` via a pointer cast, and that is sound here even though it
117 //! is UB for the `static` enum. `dynamic`'s `Tensor<T>`
118 //! (`tensor_dyn/dynamic_tensor.rs`) is `#[repr(transparent)]` with
119 //! `TensorDyn` as its only nonzero-sized field, so it has identical
120 //! layout to a bare `TensorDyn` value regardless of what `TensorDyn`'s
121 //! own fields are (a handle plus two cached facts, per
122 //! `dynamic_backend.rs`'s module docs -- not a bare pointer, but that
123 //! does not matter here: `#[repr(transparent)]` guarantees identical
124 //! layout to whatever the single field's layout is, not specifically
125 //! to a pointer's). What makes the cast safe rather than merely
126 //! well-typed is the dtype check below: it proves `T` is the caller's
127 //! actual element type before the reinterpreted reference is handed
128 //! out, exactly as `static_lens`'s `TypeId` check does for its
129 //! `downcast_ref`.
130 use super::Element;
131 use crate::{Tensor, TensorDyn};
132
133 impl TensorDyn {
134 /// Borrow this tensor as a statically-typed `Tensor<T>`, or `None`
135 /// when the element type does not match.
136 pub fn as_typed<T: Element>(&self) -> Option<&Tensor<T>> {
137 if self.dtype() != T::DTYPE {
138 return None;
139 }
140 // SAFETY: `Tensor<T>` is `#[repr(transparent)]` over `TensorDyn`
141 // (see module docs), so a `&TensorDyn` and a `&Tensor<T>` share
142 // layout; the dtype check above proves `T` matches what this
143 // handle actually holds.
144 Some(unsafe { &*(self as *const TensorDyn as *const Tensor<T>) })
145 }
146
147 /// Mutably borrow this tensor as a statically-typed `Tensor<T>`, or
148 /// `None` when the element type does not match.
149 pub fn as_typed_mut<T: Element>(&mut self) -> Option<&mut Tensor<T>> {
150 if self.dtype() != T::DTYPE {
151 return None;
152 }
153 // SAFETY: see `as_typed` above; `&mut self` gives exclusive
154 // access, preserved through the cast.
155 Some(unsafe { &mut *(self as *mut TensorDyn as *mut Tensor<T>) })
156 }
157 }
158}
159
160#[cfg(all(test, feature = "static"))]
161mod tests {
162 use crate::{DType, Tensor, TensorDyn};
163
164 #[test]
165 fn as_typed_matches_only_the_right_dtype() {
166 let t: Tensor<f32> = Tensor::new(&[4], None, None).expect("alloc");
167 let d: TensorDyn = t.into();
168 assert_eq!(d.dtype(), DType::F32);
169 assert!(d.as_typed::<f32>().is_some(), "f32 lens must open");
170 assert!(
171 d.as_typed::<u8>().is_none(),
172 "u8 lens must refuse an f32 tensor"
173 );
174 }
175
176 #[test]
177 fn as_typed_mut_is_exclusive_and_writes_through() {
178 use crate::{TensorMapTrait, TensorTrait};
179 let t: Tensor<u8> = Tensor::new(&[4], None, None).expect("alloc");
180 let mut d: TensorDyn = t.into();
181 {
182 // `map_write` takes &self and returns HostView<'_, u8>; the
183 // exclusivity here comes from `as_typed_mut`'s &mut self, not
184 // from the map. There is no HostViewMut in this codebase.
185 let lens = d.as_typed_mut::<u8>().expect("u8 lens");
186 lens.map_write().expect("map").as_mut_slice()[0] = 42;
187 }
188 let lens = d.as_typed::<u8>().expect("lens");
189 assert_eq!(lens.map_read().expect("map").as_slice()[0], 42);
190 }
191}