torsh_data/
core_framework.rs

1//! Core transform framework for data transformations
2//!
3//! This module provides the fundamental building blocks for data transformations,
4//! including core traits, combinators, and basic transform implementations.
5//!
6//! # Features
7//!
8//! - **Transform trait**: Core abstraction for data transformations
9//! - **Transform combinators**: Chain, conditional, and composition operations
10//! - **Builder pattern**: TransformBuilder trait for complex transform construction
11//! - **Extension traits**: Convenient chainable API via TransformExt
12//! - **Basic transforms**: Normalize, type conversion, and lambda transforms
13
14use torsh_core::{
15    dtype::TensorElement,
16    error::{Result, TorshError},
17};
18use torsh_tensor::Tensor;
19
20#[cfg(not(feature = "std"))]
21use alloc::{boxed::Box, string::String, vec::Vec};
22
23/// Trait for data transformations
24///
25/// This is the core abstraction for all data transformations in the ToRSh ecosystem.
26/// Implementations should be stateless where possible and thread-safe.
27pub trait Transform<T>: Send + Sync {
28    /// Output type after transformation
29    type Output;
30
31    /// Apply the transformation to a single input
32    fn transform(&self, input: T) -> Result<Self::Output>;
33
34    /// Transform multiple items in batch
35    ///
36    /// Default implementation applies transform individually, but implementations
37    /// can override this for more efficient batch processing.
38    fn transform_batch(&self, inputs: Vec<T>) -> Result<Vec<Self::Output>> {
39        inputs
40            .into_iter()
41            .map(|input| self.transform(input))
42            .collect()
43    }
44
45    /// Check if the transform is deterministic
46    ///
47    /// A deterministic transform always produces the same output for the same input.
48    /// Non-deterministic transforms include random augmentations.
49    fn is_deterministic(&self) -> bool {
50        true
51    }
52}
53
54/// Builder trait for transforms with configuration options
55pub trait TransformBuilder {
56    /// The transform type this builder creates
57    type Transform;
58
59    /// Build the configured transform
60    fn build(self) -> Self::Transform;
61}
62
63/// Macro to create simple stateless transforms
64///
65/// This macro generates a transform struct and implementation for simple cases
66/// where the transform logic can be expressed as a function.
67#[macro_export]
68macro_rules! simple_transform {
69    ($name:ident, $input:ty, $output:ty, $transform_fn:expr) => {
70        /// Auto-generated simple transform
71        #[derive(Clone, Debug, Default)]
72        pub struct $name;
73
74        impl $crate::core_framework::Transform<$input> for $name {
75            type Output = $output;
76
77            fn transform(&self, input: $input) -> $crate::core_framework::Result<Self::Output> {
78                Ok($transform_fn(input))
79            }
80        }
81    };
82
83    ($name:ident, $input:ty, $output:ty, $transform_fn:expr, deterministic = $det:literal) => {
84        /// Auto-generated simple transform with determinism setting
85        #[derive(Clone, Debug, Default)]
86        pub struct $name;
87
88        impl $crate::core_framework::Transform<$input> for $name {
89            type Output = $output;
90
91            fn transform(&self, input: $input) -> $crate::core_framework::Result<Self::Output> {
92                Ok($transform_fn(input))
93            }
94
95            fn is_deterministic(&self) -> bool {
96                $det
97            }
98        }
99    };
100}
101
102/// Extension trait for chainable transform operations
103pub trait TransformExt<T>: Transform<T> + Sized + 'static {
104    /// Chain this transform with another
105    ///
106    /// Creates a new transform that applies this transform first, then the next.
107    fn then<U>(self, next: U) -> Chain<Self, U>
108    where
109        U: Transform<Self::Output>,
110    {
111        Chain::new(self, next)
112    }
113
114    /// Apply this transform conditionally based on a predicate
115    ///
116    /// The transform is only applied if the predicate returns true for the input.
117    fn when<P>(self, predicate: P) -> Conditional<Self, P>
118    where
119        P: Fn(&T) -> bool + Send + Sync,
120    {
121        Conditional::new(self, predicate)
122    }
123
124    /// Convert to a boxed trait object for dynamic dispatch
125    fn boxed(self) -> Box<dyn Transform<T, Output = Self::Output> + Send + Sync> {
126        Box::new(self)
127    }
128}
129
130// Blanket implementation for all transforms
131impl<T, U: Transform<T> + 'static> TransformExt<T> for U {}
132
133/// Chain two transforms together sequentially
134#[derive(Debug, Clone)]
135pub struct Chain<T1, T2> {
136    first: T1,
137    second: T2,
138}
139
140impl<T1, T2> Chain<T1, T2> {
141    /// Create a new chain of transforms
142    pub fn new(first: T1, second: T2) -> Self {
143        Self { first, second }
144    }
145}
146
147impl<T, T1, T2> Transform<T> for Chain<T1, T2>
148where
149    T1: Transform<T>,
150    T2: Transform<T1::Output>,
151{
152    type Output = T2::Output;
153
154    fn transform(&self, input: T) -> Result<Self::Output> {
155        let intermediate = self.first.transform(input)?;
156        self.second.transform(intermediate)
157    }
158
159    fn is_deterministic(&self) -> bool {
160        self.first.is_deterministic() && self.second.is_deterministic()
161    }
162}
163
164/// Conditionally apply a transform based on a predicate
165#[derive(Debug, Clone)]
166pub struct Conditional<T, P> {
167    transform: T,
168    predicate: P,
169}
170
171impl<T, P> Conditional<T, P> {
172    /// Create a new conditional transform
173    pub fn new(transform: T, predicate: P) -> Self {
174        Self {
175            transform,
176            predicate,
177        }
178    }
179}
180
181impl<T, U, P> Transform<T> for Conditional<U, P>
182where
183    U: Transform<T, Output = T>,
184    P: Fn(&T) -> bool + Send + Sync,
185{
186    type Output = T;
187
188    fn transform(&self, input: T) -> Result<Self::Output> {
189        if (self.predicate)(&input) {
190            self.transform.transform(input)
191        } else {
192            Ok(input)
193        }
194    }
195
196    fn is_deterministic(&self) -> bool {
197        self.transform.is_deterministic()
198    }
199}
200
201/// Compose multiple transforms that operate on the same type
202pub struct Compose<T> {
203    transforms: Vec<Box<dyn Transform<T, Output = T> + Send + Sync>>,
204}
205
206impl<T> Compose<T> {
207    /// Create a new compose transform from a vector of transforms
208    pub fn new(transforms: Vec<Box<dyn Transform<T, Output = T> + Send + Sync>>) -> Self {
209        Self { transforms }
210    }
211
212    /// Add a transform to the composition
213    pub fn add<U>(&mut self, transform: U)
214    where
215        U: Transform<T, Output = T> + Send + Sync + 'static,
216    {
217        self.transforms.push(Box::new(transform));
218    }
219
220    /// Get the number of transforms in the composition
221    pub fn len(&self) -> usize {
222        self.transforms.len()
223    }
224
225    /// Check if the composition is empty
226    pub fn is_empty(&self) -> bool {
227        self.transforms.is_empty()
228    }
229}
230
231impl<T> Transform<T> for Compose<T> {
232    type Output = T;
233
234    fn transform(&self, mut input: T) -> Result<Self::Output> {
235        for transform in &self.transforms {
236            input = transform.transform(input)?;
237        }
238        Ok(input)
239    }
240
241    fn is_deterministic(&self) -> bool {
242        self.transforms.iter().all(|t| t.is_deterministic())
243    }
244}
245
246/// Normalize tensor values using mean and standard deviation
247#[derive(Debug, Clone)]
248pub struct Normalize<T: TensorElement> {
249    mean: Vec<T>,
250    std: Vec<T>,
251}
252
253impl<T: TensorElement> Normalize<T> {
254    /// Create a new normalize transform
255    pub fn new(mean: Vec<T>, std: Vec<T>) -> Result<Self> {
256        if mean.len() != std.len() {
257            return Err(TorshError::InvalidArgument(
258                "Mean and std vectors must have the same length".to_string(),
259            ));
260        }
261        Ok(Self { mean, std })
262    }
263}
264
265impl<
266        T: TensorElement + Copy + Default + core::ops::Sub<Output = T> + core::ops::Div<Output = T>,
267    > Transform<Tensor<T>> for Normalize<T>
268{
269    type Output = Tensor<T>;
270
271    fn transform(&self, input: Tensor<T>) -> Result<Self::Output> {
272        let num_channels = self.mean.len();
273        let ndim = input.ndim();
274
275        // Determine channel dimension:
276        // - 4D (NCHW): dim 1
277        // - 1D, 2D, 3D (CHW or CL or C): dim 0
278        let channel_dim = if ndim == 4 { 1 } else { 0 };
279
280        // For 1D tensors with 1 channel, treat the whole tensor as a single channel
281        let tensor_channels = if ndim == 0 {
282            return Err(TorshError::InvalidArgument(
283                "Normalize requires at least 1-dimensional tensor".to_string(),
284            ));
285        } else {
286            input.shape().dims()[channel_dim]
287        };
288
289        if tensor_channels != num_channels {
290            return Err(TorshError::InvalidArgument(format!(
291                "Tensor has {} channels at dim {} but Normalize has {} channels in mean/std",
292                tensor_channels, channel_dim, num_channels
293            )));
294        }
295
296        // Validate std is nonzero for all channels
297        for (c, std_val) in self.std.iter().enumerate() {
298            if std_val.is_zero() {
299                return Err(TorshError::InvalidArgument(format!(
300                    "std[{c}] is zero, which would cause division by zero in Normalize"
301                )));
302            }
303        }
304
305        // Process each channel: (x - mean[c]) / std[c]
306        let mut channel_slices: Vec<Tensor<T>> = Vec::with_capacity(num_channels);
307        for c in 0..num_channels {
308            let channel_slice = input.slice_tensor(channel_dim, c, c + 1)?;
309            let centered = channel_slice.sub_scalar(self.mean[c])?;
310            let normalized = centered.div_scalar(self.std[c])?;
311            channel_slices.push(normalized);
312        }
313
314        let channel_refs: Vec<&Tensor<T>> = channel_slices.iter().collect();
315        Tensor::cat(&channel_refs, channel_dim as i32)
316    }
317}
318
319/// Convert tensor from one type to another
320#[derive(Debug, Clone)]
321pub struct ToType<From, To> {
322    _phantom: core::marker::PhantomData<(From, To)>,
323}
324
325impl<From, To> Default for ToType<From, To> {
326    fn default() -> Self {
327        Self::new()
328    }
329}
330
331impl<From, To> ToType<From, To> {
332    /// Create a new type conversion transform
333    pub fn new() -> Self {
334        Self {
335            _phantom: core::marker::PhantomData,
336        }
337    }
338}
339
340impl<From: TensorElement, To: TensorElement> Transform<Tensor<From>> for ToType<From, To> {
341    type Output = Tensor<To>;
342
343    fn transform(&self, _input: Tensor<From>) -> Result<Self::Output> {
344        // Placeholder implementation - real type conversion would require tensor operations
345        // For now, create a new tensor with the target type (this is a simplification)
346        // NOTE: tracing disabled (not a dependency)
347        // tracing::debug!(
348        //     "Type conversion from {} to {} requested",
349        //     core::any::type_name::<From>(),
350        //     core::any::type_name::<To>()
351        // );
352
353        // In a real implementation, this would convert the tensor data
354        // For now, we return an error as this requires complex tensor operations
355        Err(TorshError::InvalidArgument(
356            "Type conversion not yet implemented".to_string(),
357        ))
358    }
359}
360
361/// Apply a custom function as a transform
362#[derive(Debug)]
363pub struct Lambda<F> {
364    func: F,
365}
366
367impl<F> Lambda<F> {
368    /// Create a new lambda transform
369    pub fn new(func: F) -> Self {
370        Self { func }
371    }
372}
373
374impl<T, O, F> Transform<T> for Lambda<F>
375where
376    F: Fn(T) -> Result<O> + Send + Sync,
377{
378    type Output = O;
379
380    fn transform(&self, input: T) -> Result<Self::Output> {
381        (self.func)(input)
382    }
383
384    fn is_deterministic(&self) -> bool {
385        // Lambda functions are assumed to be deterministic unless specified otherwise
386        true
387    }
388}
389
390/// Convenience function to create a normalize transform
391pub fn normalize<T: TensorElement>(mean: Vec<T>, std: Vec<T>) -> Result<Normalize<T>> {
392    Normalize::new(mean, std)
393}
394
395/// Convenience function to create a type conversion transform
396pub fn to_type<From: TensorElement, To: TensorElement>() -> ToType<From, To> {
397    ToType::new()
398}
399
400/// Convenience function to create a lambda transform
401pub fn lambda<F>(func: F) -> Lambda<F> {
402    Lambda::new(func)
403}
404
405/// Convenience function to create a composition transform
406pub fn compose<T>(transforms: Vec<Box<dyn Transform<T, Output = T> + Send + Sync>>) -> Compose<T> {
407    Compose::new(transforms)
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    // Mock tensor for testing
415    #[allow(dead_code)]
416    fn mock_tensor() -> Tensor<f32> {
417        Tensor::from_data(
418            vec![1.0f32, 2.0, 3.0, 4.0],
419            vec![2, 2],
420            torsh_core::device::DeviceType::Cpu,
421        )
422        .unwrap()
423    }
424
425    #[test]
426    fn test_chain_transform() {
427        let lambda1 = lambda(|x: i32| Ok(x * 2));
428        let lambda2 = lambda(|x: i32| Ok(x + 1));
429
430        let chained = lambda1.then(lambda2);
431        let result = chained.transform(5).unwrap();
432        assert_eq!(result, 11); // (5 * 2) + 1 = 11
433    }
434
435    #[test]
436    fn test_conditional_transform() {
437        let double = lambda(|x: i32| Ok(x * 2));
438        let conditional = double.when(|&x| x > 5);
439
440        assert_eq!(conditional.transform(3).unwrap(), 3); // Not applied
441        assert_eq!(conditional.transform(7).unwrap(), 14); // Applied
442    }
443
444    #[test]
445    fn test_compose_transform() {
446        let lambda1 = lambda(|x: i32| Ok(x + 1));
447        let lambda2 = lambda(|x: i32| Ok(x * 2));
448
449        let mut composition = Compose::new(vec![]);
450        composition.add(lambda1);
451        composition.add(lambda2);
452
453        let result = composition.transform(5).unwrap();
454        assert_eq!(result, 12); // ((5 + 1) * 2) = 12
455    }
456
457    #[test]
458    fn test_normalize_creation() {
459        let mean = vec![0.485f32, 0.456, 0.406];
460        let std = vec![0.229f32, 0.224, 0.225];
461
462        let normalize_transform = normalize(mean, std);
463        assert!(normalize_transform.is_ok());
464    }
465
466    #[test]
467    fn test_normalize_invalid_dimensions() {
468        let mean = vec![0.485f32, 0.456];
469        let std = vec![0.229f32, 0.224, 0.225];
470
471        let normalize_transform = normalize(mean, std);
472        assert!(normalize_transform.is_err());
473    }
474
475    #[test]
476    fn test_determinism() {
477        let deterministic = lambda(|x: i32| Ok(x + 1));
478        assert!(deterministic.is_deterministic());
479
480        let chain = deterministic.then(lambda(|x: i32| Ok(x * 2)));
481        assert!(chain.is_deterministic());
482    }
483
484    #[test]
485    fn test_normalize_transform_3channel_chw() {
486        use torsh_core::device::DeviceType;
487
488        // 3-channel CHW tensor: shape [3, 2, 2]
489        // Channel 0: all 1.0, Channel 1: all 3.0, Channel 2: all 5.0
490        let data = vec![
491            1.0f32, 1.0, 1.0, 1.0, // channel 0
492            3.0f32, 3.0, 3.0, 3.0, // channel 1
493            5.0f32, 5.0, 5.0, 5.0, // channel 2
494        ];
495        let input = Tensor::from_data(data, vec![3, 2, 2], DeviceType::Cpu).unwrap();
496
497        let mean = vec![0.0f32, 1.0, 2.0];
498        let std = vec![1.0f32, 2.0, 1.0];
499        let norm = Normalize::new(mean, std).unwrap();
500
501        let output = norm.transform(input).unwrap();
502        assert_eq!(output.shape().dims(), &[3, 2, 2]);
503
504        let out_data = output.data().unwrap();
505        // Channel 0: (1.0 - 0.0) / 1.0 = 1.0
506        assert!(
507            (out_data[0] - 1.0f32).abs() < 1e-5,
508            "ch0 expected 1.0, got {}",
509            out_data[0]
510        );
511        assert!(
512            (out_data[1] - 1.0f32).abs() < 1e-5,
513            "ch0 expected 1.0, got {}",
514            out_data[1]
515        );
516        // Channel 1: (3.0 - 1.0) / 2.0 = 1.0
517        assert!(
518            (out_data[4] - 1.0f32).abs() < 1e-5,
519            "ch1 expected 1.0, got {}",
520            out_data[4]
521        );
522        // Channel 2: (5.0 - 2.0) / 1.0 = 3.0
523        assert!(
524            (out_data[8] - 3.0f32).abs() < 1e-5,
525            "ch2 expected 3.0, got {}",
526            out_data[8]
527        );
528    }
529
530    #[test]
531    fn test_normalize_transform_channel_mismatch() {
532        use torsh_core::device::DeviceType;
533
534        let input = Tensor::from_data(
535            vec![1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0],
536            vec![3, 2],
537            DeviceType::Cpu,
538        )
539        .unwrap();
540
541        // mean/std have 2 channels but tensor has 3
542        let mean = vec![0.0f32, 1.0];
543        let std = vec![1.0f32, 2.0];
544        let norm = Normalize::new(mean, std).unwrap();
545
546        let result = norm.transform(input);
547        assert!(result.is_err(), "Expected error due to channel mismatch");
548    }
549
550    #[test]
551    fn test_normalize_transform_zero_std() {
552        use torsh_core::device::DeviceType;
553
554        let input =
555            Tensor::from_data(vec![1.0f32, 2.0, 3.0, 4.0], vec![2, 2], DeviceType::Cpu).unwrap();
556
557        let mean = vec![0.0f32, 1.0];
558        let std = vec![1.0f32, 0.0]; // std[1] = 0, invalid
559        let norm = Normalize::new(mean, std).unwrap();
560
561        let result = norm.transform(input);
562        assert!(result.is_err(), "Expected error due to zero std");
563    }
564}
torsh_data/core_framework.rs

torsh_data/
core_framework.rs
