mod cpu_kernel;
#[cfg(feature = "cuda")]
mod cuda_kernel;
use super::ops::*;
use crate::{shapes::*, tensor::*};
#[repr(C)]
#[derive(Debug, Default, Clone, Copy)]
pub struct BinaryMulKernelOp;
#[repr(C)]
#[derive(Debug, Clone, Copy)]
pub struct ScalarMulKernelOp<E> {
scalar: E,
}
pub fn mul<S: Shape, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Default>(
lhs: Tensor<S, E, D, T>,
rhs: Tensor<S, E, D, R>,
) -> Tensor<S, E, D, T>
where
D: BinaryKernel<BinaryMulKernelOp, E>,
{
lhs * rhs
}
pub trait TryMul<Rhs = Self>: HasErr {
fn try_mul(self, rhs: Rhs) -> Result<Self, Self::Err>;
}
impl<S: Shape, E: Dtype, D: BinaryKernel<BinaryMulKernelOp, E>, LhsTape: Tape<E, D>, R>
TryMul<Tensor<S, E, D, R>> for Tensor<S, E, D, LhsTape>
where
LhsTape: Merge<R>,
{
fn try_mul(self, rhs: Tensor<S, E, D, R>) -> Result<Self, Self::Err> {
try_binary_op(BinaryMulKernelOp, self, rhs)
}
}
impl<S: Shape, E: Dtype, D: UnaryKernel<ScalarMulKernelOp<E>, E>, T: Tape<E, D>> TryMul<E>
for Tensor<S, E, D, T>
{
fn try_mul(self, rhs: E) -> Result<Self, Self::Err> {
try_unary_op(ScalarMulKernelOp { scalar: rhs }, self)
}
}
#[cfg(feature = "f16")]
impl<S: Shape, D: UnaryKernel<ScalarMulKernelOp<half::f16>, half::f16>, T: Tape<half::f16, D>>
TryMul<f32> for Tensor<S, half::f16, D, T>
{
fn try_mul(self, rhs: f32) -> Result<Self, Self::Err> {
let scalar = half::f16::from_f32(rhs);
try_unary_op(ScalarMulKernelOp { scalar }, self)
}
}
#[cfg(feature = "f16")]
impl<
S: Shape,
D: UnaryKernel<
ScalarMulKernelOp<crate::dtypes::AMP<half::f16>>,
crate::dtypes::AMP<half::f16>,
>,
T: Tape<crate::dtypes::AMP<half::f16>, D>,
> TryMul<f32> for Tensor<S, crate::dtypes::AMP<half::f16>, D, T>
{
fn try_mul(self, rhs: f32) -> Result<Self, Self::Err> {
let scalar = crate::dtypes::AMP(half::f16::from_f32(rhs));
try_unary_op(ScalarMulKernelOp { scalar }, self)
}
}
impl<S: Shape, E: Dtype, D: Storage<E>, LhsTape: Tape<E, D>, Rhs> std::ops::Mul<Rhs>
for Tensor<S, E, D, LhsTape>
where
Self: TryMul<Rhs>,
{
type Output = Self;
fn mul(self, rhs: Rhs) -> Self::Output {
self.try_mul(rhs).unwrap()
}
}
#[cfg(test)]
mod tests {
use crate::{tensor::*, tensor_ops::*, tests::*};
#[test]
fn test_mul_0d() {
let dev: TestDevice = Default::default();
let a = dev.tensor(2.0).to_dtype::<TestDtype>();
let b = dev.tensor(3.0).to_dtype::<TestDtype>();
let r = a.leaky_trace() * b.clone();
assert_close_to_literal!(r, 6.0);
let g = r.backward();
assert_close_to_literal!(g.get(&a), 3.0);
assert_close_to_literal!(g.get(&b), 2.0);
}
#[test]
fn test_mul_1d() {
let dev: TestDevice = Default::default();
let a = dev.tensor([1.0, 2.0, 3.0]).to_dtype::<TestDtype>();
let b = dev.tensor([1.0, -1.0, 0.0]).to_dtype::<TestDtype>();
let r = a.leaky_trace() * b.clone();
assert_close_to_literal!(r, [1.0, -2.0, 0.0]);
let g = r.mean().backward();
assert_close_to_literal!(g.get(&a), [1.0 / 3.0, -1.0 / 3.0, 0.0]);
assert_close_to_literal!(g.get(&b), [1.0 / 3.0, 2.0 / 3.0, 1.0]);
}
#[test]
fn test_mul_2d() {
let dev: TestDevice = Default::default();
let a = dev
.tensor([[0.6570, 0.1708, 0.1500], [0.5658, 0.7010, 0.8342]])
.to_dtype::<TestDtype>();
let b = dev
.tensor([[0.5199, 0.3844, 0.3759], [0.8259, 0.3682, 0.0388]])
.to_dtype::<TestDtype>();
let r = a.leaky_trace() * b.clone();
assert_close_to_literal!(
r,
[
[0.3415743, 0.06565552, 0.056385003],
[0.46729425, 0.2581082, 0.03236696],
]
);
let g = r.mean().backward();
assert_close_to_literal!(
g.get(&a),
[
[0.08665001, 0.06406667, 0.06265],
[0.13765001, 0.06136667, 0.006466667],
]
);
assert_close_to_literal!(
g.get(&b),
[
[0.109500006, 0.028466668, 0.025000002],
[0.0943, 0.11683333, 0.13903335],
]
);
}
#[test]
fn test_scalar_mul_0d() {
let dev: TestDevice = Default::default();
let x = dev.tensor(1.0).to_dtype::<TestDtype>();
let r = x.leaky_trace() * 0.5;
assert_close_to_literal!(r, 0.5);
let g = r.exp().backward();
assert_close_to_literal!(g.get(&x), 0.8243606);
}
#[test]
fn test_scalar_mul_1d() {
let dev: TestDevice = Default::default();
let x = dev.tensor([0.0, 1.0, 2.0]).to_dtype::<TestDtype>();
let r = x.leaky_trace() * 0.5;
assert_close_to_literal!(r, [0.0, 0.5, 1.0]);
let g = r.exp().sum().backward();
assert_close_to_literal!(g.get(&x), [0.5, 0.8243606, 1.3591409]);
}
#[test]
fn test_scalar_mul_2d() {
let dev: TestDevice = Default::default();
let x = dev.tensor([[1.0; 2]; 3]).to_dtype::<TestDtype>();
let r = x.leaky_trace() * 0.5;
assert_close_to_literal!(r, [[0.5; 2]; 3]);
let g = r.exp().sum().backward();
assert_close_to_literal!(g.get(&x), [[0.8243606; 2]; 3]);
}
}