torsh-functional 0.1.2

Functional programming utilities for ToRSh tensors
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

//! Common types and utilities for loss functions
//!
//! This module provides shared types and helper functions used across
//! different loss function categories.

use torsh_core::Result as TorshResult;
use torsh_tensor::Tensor;

/// Reduction type for loss functions
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ReductionType {
    /// No reduction applied
    None,
    /// Mean reduction
    Mean,
    /// Sum reduction
    Sum,
}

impl ReductionType {
    /// Apply the reduction to a tensor
    pub fn apply(&self, tensor: Tensor) -> TorshResult<Tensor> {
        match self {
            Self::None => Ok(tensor),
            Self::Mean => tensor.mean(None, false),
            Self::Sum => tensor.sum(),
        }
    }
}