hpt 0.1.3

High Performance Tensor (HPT) - A fast, efficient, and user-friendly tensor computation library for Rust
Documentation
use std::borrow::{Borrow, BorrowMut};

use crate::Tensor;
use hpt_allocator::traits::{Allocator, AllocatorOutputRetrive};
use hpt_allocator::Cpu;
use hpt_common::error::base::TensorError;
use hpt_traits::ops::binary::Gemm;
use hpt_traits::tensor::CommonBounds;
use hpt_types::{into_scalar::Cast, type_promote::NormalOut};

type GemmOutput<A, B, const DEVICE: usize, A2> =
    Tensor<<A as NormalOut<B>>::Output, Cpu, DEVICE, A2>;

impl<A, B, A2, const DEVICE: usize> Gemm<Tensor<B, Cpu, DEVICE, A2>> for Tensor<A, Cpu, DEVICE, A2>
where
    A: CommonBounds + NormalOut<B> + Cast<<A as NormalOut<B>>::Output>,
    B: CommonBounds + Cast<<A as NormalOut<B>>::Output>,
    <A as NormalOut<B>>::Output: CommonBounds,
    A2: Allocator,
    A2::Output: AllocatorOutputRetrive,
{
    type Output = GemmOutput<A, B, DEVICE, A2>;

    type OutputMeta = <A as NormalOut<B>>::Output;

    type InplaceOutput = GemmOutput<A, B, DEVICE, A2>;

    fn gemm(
        &self,
        rhs: Tensor<B, Cpu, DEVICE, A2>,
        alpha: Self::OutputMeta,
        beta: Self::OutputMeta,
        conj_dst: bool,
        conj_lhs: bool,
        conj_rhs: bool,
    ) -> Result<Self::Output, TensorError> {
        Ok(self
            .inner
            .gemm(
                rhs.inner.as_ref(),
                alpha,
                beta,
                conj_dst,
                conj_lhs,
                conj_rhs,
            )?
            .into())
    }
    fn gemm_<U>(
        &self,
        rhs: Tensor<B, Cpu, DEVICE, A2>,
        alpha: Self::OutputMeta,
        beta: Self::OutputMeta,
        conj_dst: bool,
        conj_lhs: bool,
        conj_rhs: bool,
        out: U,
    ) -> Result<Self::Output, TensorError>
    where
        U: Borrow<Self::InplaceOutput> + BorrowMut<Self::InplaceOutput>,
    {
        let out = out.borrow().inner.as_ref().clone();
        Ok(self
            .inner
            .gemm_(
                rhs.inner.as_ref(),
                alpha,
                beta,
                conj_dst,
                conj_lhs,
                conj_rhs,
                out,
            )?
            .into())
    }
}

impl<A, B, A2, const DEVICE: usize> Gemm<&Tensor<B, Cpu, DEVICE, A2>> for Tensor<A, Cpu, DEVICE, A2>
where
    A: CommonBounds + NormalOut<B> + Cast<<A as NormalOut<B>>::Output>,
    B: CommonBounds + Cast<<A as NormalOut<B>>::Output>,
    <A as NormalOut<B>>::Output: CommonBounds,
    A2: Allocator,
    A2::Output: AllocatorOutputRetrive,
{
    type Output = GemmOutput<A, B, DEVICE, A2>;

    type OutputMeta = <A as NormalOut<B>>::Output;

    type InplaceOutput = GemmOutput<A, B, DEVICE, A2>;

    fn gemm(
        &self,
        rhs: &Tensor<B, Cpu, DEVICE, A2>,
        alpha: Self::OutputMeta,
        beta: Self::OutputMeta,
        conj_dst: bool,
        conj_lhs: bool,
        conj_rhs: bool,
    ) -> Result<Self::Output, TensorError> {
        Ok(self
            .inner
            .gemm(
                rhs.inner.as_ref(),
                alpha,
                beta,
                conj_dst,
                conj_lhs,
                conj_rhs,
            )?
            .into())
    }

    fn gemm_<U>(
        &self,
        rhs: &Tensor<B, Cpu, DEVICE, A2>,
        alpha: Self::OutputMeta,
        beta: Self::OutputMeta,
        conj_dst: bool,
        conj_lhs: bool,
        conj_rhs: bool,
        out: U,
    ) -> Result<Self::Output, TensorError>
    where
        U: Borrow<Self::InplaceOutput> + BorrowMut<Self::InplaceOutput>,
    {
        let out = out.borrow().inner.as_ref().clone();
        Ok(self
            .inner
            .gemm_(
                rhs.inner.as_ref(),
                alpha,
                beta,
                conj_dst,
                conj_lhs,
                conj_rhs,
                out,
            )?
            .into())
    }
}