From 7df12c0dd913a4257934fe9a7233ade30382bcd3 Mon Sep 17 00:00:00 2001 From: Corey Lowman Date: Wed, 25 Oct 2023 12:52:01 -0400 Subject: [PATCH] [Breaking] Combining separate device errors into single `dfdx::tensor::Error` enum (#875) * [Breaking] Adding single Error enum * Fixing example * Fixing cuda kernels * Fixing no-std --- dfdx-core/src/nn_traits/mod.rs | 56 ++++--------- dfdx-core/src/nn_traits/tuples.rs | 21 +++-- dfdx-core/src/nn_traits/vecs.rs | 22 +++-- dfdx-core/src/tensor/cpu/allocate.rs | 39 ++++----- dfdx-core/src/tensor/cpu/device.rs | 36 ++------ dfdx-core/src/tensor/cpu/mod.rs | 2 +- dfdx-core/src/tensor/cuda/allocate.rs | 26 +++--- dfdx-core/src/tensor/cuda/device.rs | 61 ++++---------- dfdx-core/src/tensor/cuda/mod.rs | 2 +- dfdx-core/src/tensor/error.rs | 28 +++++++ dfdx-core/src/tensor/ghost.rs | 6 +- dfdx-core/src/tensor/gradients.rs | 20 ++--- dfdx-core/src/tensor/mod.rs | 8 +- dfdx-core/src/tensor/storage_traits.rs | 82 +++++++++---------- dfdx-core/src/tensor/tensor_impls.rs | 10 +-- dfdx-core/src/tensor/tensorlike.rs | 8 +- dfdx-core/src/tensor_ops/abs/mod.rs | 2 +- dfdx-core/src/tensor_ops/accurate_gelu/mod.rs | 2 +- dfdx-core/src/tensor_ops/adam/cpu_kernel.rs | 6 +- dfdx-core/src/tensor_ops/adam/cuda_kernel.rs | 4 +- dfdx-core/src/tensor_ops/adam/mod.rs | 6 +- dfdx-core/src/tensor_ops/add/mod.rs | 10 +-- .../attention_reshape/cpu_kernel.rs | 2 +- .../attention_reshape/cuda_kernel.rs | 2 +- .../src/tensor_ops/attention_reshape/mod.rs | 6 +- dfdx-core/src/tensor_ops/axpy/cpu_kernel.rs | 13 ++- dfdx-core/src/tensor_ops/axpy/cuda_kernel.rs | 10 +-- dfdx-core/src/tensor_ops/axpy/mod.rs | 7 +- dfdx-core/src/tensor_ops/bce/mod.rs | 5 +- .../src/tensor_ops/boolean/cpu_kernels.rs | 15 ++-- .../src/tensor_ops/boolean/cuda_kernels.rs | 15 ++-- dfdx-core/src/tensor_ops/boolean/mod.rs | 19 ++--- dfdx-core/src/tensor_ops/broadcast_to.rs | 8 +- dfdx-core/src/tensor_ops/choose/cpu_kernel.rs | 6 +- .../src/tensor_ops/choose/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/choose/mod.rs | 12 +-- dfdx-core/src/tensor_ops/clamp/mod.rs | 6 +- dfdx-core/src/tensor_ops/cmp/cpu_kernels.rs | 6 +- dfdx-core/src/tensor_ops/cmp/cuda_kernels.rs | 6 +- dfdx-core/src/tensor_ops/cmp/mod.rs | 24 +++--- dfdx-core/src/tensor_ops/concat/cpu_kernel.rs | 6 +- .../src/tensor_ops/concat/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/concat/mod.rs | 10 +-- .../src/tensor_ops/concat_along/cpu_kernel.rs | 4 +- .../tensor_ops/concat_along/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/concat_along/mod.rs | 16 ++-- dfdx-core/src/tensor_ops/conv1d/cpu_kernel.rs | 10 +-- .../src/tensor_ops/conv1d/cuda_kernel.rs | 8 +- dfdx-core/src/tensor_ops/conv1d/mod.rs | 23 ++---- dfdx-core/src/tensor_ops/conv2d/cpu_kernel.rs | 10 +-- .../src/tensor_ops/conv2d/cuda_kernel.rs | 8 +- .../src/tensor_ops/conv2d/cudnn_kernel.rs | 8 +- dfdx-core/src/tensor_ops/conv2d/mod.rs | 21 ++--- .../src/tensor_ops/convtrans2d/cpu_kernel.rs | 12 +-- .../src/tensor_ops/convtrans2d/cuda_kernel.rs | 8 +- dfdx-core/src/tensor_ops/convtrans2d/mod.rs | 21 ++--- dfdx-core/src/tensor_ops/cos/mod.rs | 2 +- dfdx-core/src/tensor_ops/div/mod.rs | 8 +- .../src/tensor_ops/dropout/cpu_kernel.rs | 6 +- .../src/tensor_ops/dropout/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/dropout/mod.rs | 11 +-- dfdx-core/src/tensor_ops/exp/mod.rs | 2 +- dfdx-core/src/tensor_ops/fast_gelu/mod.rs | 4 +- dfdx-core/src/tensor_ops/huber_error/mod.rs | 2 +- dfdx-core/src/tensor_ops/ln/mod.rs | 2 +- dfdx-core/src/tensor_ops/log_softmax.rs | 2 +- dfdx-core/src/tensor_ops/logsumexp_to.rs | 6 +- dfdx-core/src/tensor_ops/matmul/cpu_kernel.rs | 18 ++-- .../src/tensor_ops/matmul/cuda_kernel.rs | 18 ++-- dfdx-core/src/tensor_ops/matmul/mod.rs | 42 +++++----- dfdx-core/src/tensor_ops/max_to/cpu_kernel.rs | 6 +- .../src/tensor_ops/max_to/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/max_to/mod.rs | 10 +-- dfdx-core/src/tensor_ops/maximum/mod.rs | 5 +- dfdx-core/src/tensor_ops/mean_to.rs | 6 +- dfdx-core/src/tensor_ops/min_to/cpu_kernel.rs | 6 +- .../src/tensor_ops/min_to/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/min_to/mod.rs | 10 +-- dfdx-core/src/tensor_ops/minimum/mod.rs | 5 +- dfdx-core/src/tensor_ops/mul/mod.rs | 8 +- dfdx-core/src/tensor_ops/nans_to/mod.rs | 2 +- dfdx-core/src/tensor_ops/negate/mod.rs | 2 +- dfdx-core/src/tensor_ops/normalize.rs | 7 +- dfdx-core/src/tensor_ops/permute_to.rs | 6 +- dfdx-core/src/tensor_ops/pool2d/cpu_kernel.rs | 6 +- .../src/tensor_ops/pool2d/cuda_kernel.rs | 8 +- dfdx-core/src/tensor_ops/pool2d/mod.rs | 21 ++--- dfdx-core/src/tensor_ops/pow/cuda_kernel.rs | 4 +- dfdx-core/src/tensor_ops/pow/mod.rs | 4 +- dfdx-core/src/tensor_ops/prelu.rs | 8 +- dfdx-core/src/tensor_ops/realize_to.rs | 2 +- dfdx-core/src/tensor_ops/recip/mod.rs | 2 +- dfdx-core/src/tensor_ops/relu/mod.rs | 2 +- .../src/tensor_ops/reshape_to/cpu_kernel.rs | 6 +- .../src/tensor_ops/reshape_to/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/reshape_to/mod.rs | 14 ++-- .../src/tensor_ops/rmsprop/cpu_kernel.rs | 6 +- .../src/tensor_ops/rmsprop/cuda_kernel.rs | 4 +- dfdx-core/src/tensor_ops/rmsprop/mod.rs | 4 +- dfdx-core/src/tensor_ops/roll/cpu_kernel.rs | 4 +- dfdx-core/src/tensor_ops/roll/cuda_kernel.rs | 4 +- dfdx-core/src/tensor_ops/roll/mod.rs | 13 +-- .../select_and_gather/cpu_kernel.rs | 10 +-- .../select_and_gather/cuda_kernel.rs | 10 +-- .../src/tensor_ops/select_and_gather/mod.rs | 20 ++--- dfdx-core/src/tensor_ops/sgd/cpu_kernel.rs | 6 +- dfdx-core/src/tensor_ops/sgd/cuda_kernel.rs | 4 +- dfdx-core/src/tensor_ops/sgd/mod.rs | 6 +- dfdx-core/src/tensor_ops/sigmoid/mod.rs | 2 +- dfdx-core/src/tensor_ops/sin/mod.rs | 2 +- dfdx-core/src/tensor_ops/slice/cpu_kernel.rs | 4 +- dfdx-core/src/tensor_ops/slice/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/slice/mod.rs | 9 +- dfdx-core/src/tensor_ops/softmax.rs | 2 +- dfdx-core/src/tensor_ops/sqrt/mod.rs | 2 +- dfdx-core/src/tensor_ops/square/mod.rs | 2 +- dfdx-core/src/tensor_ops/stack/cpu_kernel.rs | 6 +- dfdx-core/src/tensor_ops/stack/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/stack/mod.rs | 24 ++---- dfdx-core/src/tensor_ops/stddev_to.rs | 6 +- dfdx-core/src/tensor_ops/sub/mod.rs | 8 +- dfdx-core/src/tensor_ops/sum_to/cpu_kernel.rs | 10 +-- .../src/tensor_ops/sum_to/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/sum_to/mod.rs | 10 +-- dfdx-core/src/tensor_ops/tanh/mod.rs | 2 +- .../src/tensor_ops/to_dtype/cpu_kernel.rs | 4 +- .../src/tensor_ops/to_dtype/cuda_kernel.rs | 4 +- dfdx-core/src/tensor_ops/to_dtype/mod.rs | 6 +- dfdx-core/src/tensor_ops/tri.rs | 18 ++-- .../src/tensor_ops/upscale2d/cpu_kernel.rs | 10 +-- .../src/tensor_ops/upscale2d/cuda_kernel.rs | 6 +- dfdx-core/src/tensor_ops/upscale2d/mod.rs | 20 ++--- .../src/tensor_ops/utilities/backward.rs | 8 +- .../src/tensor_ops/utilities/cpu_kernels.rs | 10 +-- .../src/tensor_ops/utilities/cuda_kernels.rs | 8 +- dfdx-core/src/tensor_ops/utilities/ops.rs | 14 ++-- dfdx-core/src/tensor_ops/var_to.rs | 6 +- dfdx-derives/src/lib.rs | 42 ++++------ dfdx/examples/advanced-train-loop.rs | 6 +- dfdx/src/nn/layers/abs.rs | 3 +- dfdx/src/nn/layers/add_into.rs | 11 ++- dfdx/src/nn/layers/batch_norm1d.rs | 21 ++--- dfdx/src/nn/layers/batch_norm2d.rs | 24 ++---- dfdx/src/nn/layers/bias1d.rs | 16 ++-- dfdx/src/nn/layers/bias2d.rs | 10 +-- dfdx/src/nn/layers/conv1d.rs | 13 +-- dfdx/src/nn/layers/conv2d.rs | 11 +-- dfdx/src/nn/layers/conv_trans2d.rs | 11 +-- dfdx/src/nn/layers/cos.rs | 3 +- dfdx/src/nn/layers/dropout.rs | 10 +-- dfdx/src/nn/layers/embedding.rs | 13 +-- dfdx/src/nn/layers/exp.rs | 3 +- dfdx/src/nn/layers/flatten2d.rs | 9 +- dfdx/src/nn/layers/gelu.rs | 6 +- dfdx/src/nn/layers/generalized_add.rs | 12 ++- dfdx/src/nn/layers/generalized_mul.rs | 12 ++- dfdx/src/nn/layers/layer_norm1d.rs | 16 ++-- dfdx/src/nn/layers/leaky_relu.rs | 3 +- dfdx/src/nn/layers/linear.rs | 12 ++- dfdx/src/nn/layers/ln.rs | 3 +- dfdx/src/nn/layers/log_softmax.rs | 3 +- dfdx/src/nn/layers/matmul.rs | 9 +- dfdx/src/nn/layers/multi_head_attention.rs | 11 +-- dfdx/src/nn/layers/pool_2d_avg.rs | 3 +- dfdx/src/nn/layers/pool_2d_max.rs | 4 +- dfdx/src/nn/layers/pool_2d_min.rs | 4 +- dfdx/src/nn/layers/pool_global_avg.rs | 14 ++-- dfdx/src/nn/layers/pool_global_max.rs | 14 ++-- dfdx/src/nn/layers/pool_global_min.rs | 14 ++-- dfdx/src/nn/layers/prelu.rs | 8 +- dfdx/src/nn/layers/prelu1d.rs | 20 ++--- dfdx/src/nn/layers/relu.rs | 3 +- dfdx/src/nn/layers/reshape.rs | 3 +- dfdx/src/nn/layers/residual_add.rs | 10 +-- dfdx/src/nn/layers/residual_mul.rs | 10 +-- dfdx/src/nn/layers/sigmoid.rs | 3 +- dfdx/src/nn/layers/sin.rs | 3 +- dfdx/src/nn/layers/softmax.rs | 3 +- dfdx/src/nn/layers/split_into.rs | 11 +-- dfdx/src/nn/layers/sqrt.rs | 3 +- dfdx/src/nn/layers/square.rs | 3 +- dfdx/src/nn/layers/tanh.rs | 3 +- dfdx/src/nn/layers/transformer.rs | 21 ++--- dfdx/src/nn/layers/upscale2d.rs | 12 +-- dfdx/src/nn/optim/adam.rs | 18 ++-- dfdx/src/nn/optim/mod.rs | 2 +- dfdx/src/nn/optim/rmsprop.rs | 18 ++-- dfdx/src/nn/optim/sgd.rs | 2 +- 188 files changed, 844 insertions(+), 1069 deletions(-) create mode 100644 dfdx-core/src/tensor/error.rs diff --git a/dfdx-core/src/nn_traits/mod.rs b/dfdx-core/src/nn_traits/mod.rs index 204d82dff..20c55da23 100644 --- a/dfdx-core/src/nn_traits/mod.rs +++ b/dfdx-core/src/nn_traits/mod.rs @@ -3,17 +3,16 @@ mod vecs; use std::vec::Vec; -use crate::prelude::{Device, Dtype, Gradients, Shape, Tensor, UniqueId}; +use crate::prelude::{Device, Dtype, Error, Gradients, Shape, Tensor, UniqueId}; /// Mutable & Immutable forward of `Input` that produces [Module::Output]. pub trait Module { /// The type that this unit produces given `Input`. type Output; - type Error: std::fmt::Debug; - fn try_forward(&self, x: X) -> Result; + fn try_forward(&self, x: X) -> Result; - fn try_forward_mut(&mut self, x: X) -> Result { + fn try_forward_mut(&mut self, x: X) -> Result { self.try_forward(x) } @@ -26,27 +25,6 @@ pub trait Module { } } -/// An error indicating that a parameter was not used in gradient -/// computation, and was therefore not present in [Gradients] -/// during an update. -#[derive(Debug)] -pub enum OptimizerUpdateError { - UnusedTensors(Vec), - DeviceError(Err), -} - -impl std::fmt::Display for OptimizerUpdateError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::UnusedTensors(unused) => write!(f, "Unused tensors: {unused:?}"), - Self::DeviceError(err) => write!(f, "{err}"), - } - } -} - -#[cfg(feature = "std")] -impl std::error::Error for OptimizerUpdateError {} - /// Something that can update both tensors and a [UpdateParams]. At minimum [Optimizer::update_tensor()] must be implemented. pub trait Optimizer>: Sized { fn update_tensor( @@ -54,24 +32,18 @@ pub trait Optimizer>: Sized { t: &mut Tensor, gradients: &Gradients, missing_tensors: &mut Vec, - ) -> Result<(), D::Err>; + ) -> Result<(), Error>; - fn update( - &mut self, - module: &mut M, - gradients: &Gradients, - ) -> Result<(), OptimizerUpdateError> + fn update(&mut self, module: &mut M, gradients: &Gradients) -> Result<(), Error> where M: UpdateParams, { let mut missing_tensors = Vec::new(); - module - .try_update_params(self, gradients, &mut missing_tensors) - .map_err(OptimizerUpdateError::DeviceError)?; + module.try_update_params(self, gradients, &mut missing_tensors)?; if missing_tensors.is_empty() { Ok(()) } else { - Err(OptimizerUpdateError::UnusedTensors(missing_tensors)) + Err(Error::UnusedTensors(missing_tensors)) } } } @@ -82,7 +54,7 @@ pub trait BuildOnDevice>: Clone { fn build_on_device(&self, device: &D) -> Self::Built { self.try_build_on_device(device).unwrap() } - fn try_build_on_device(&self, device: &D) -> Result; + fn try_build_on_device(&self, device: &D) -> Result; } /// Something that can have all of its parameters reset to a specific state (may be random or not random). @@ -90,7 +62,7 @@ pub trait ResetParams> { fn reset_params(&mut self) { self.try_reset_params().unwrap() } - fn try_reset_params(&mut self) -> Result<(), D::Err>; + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error>; } /// Something that can have it's params updated with an [Optimizer] and a set of [Gradients]. @@ -109,7 +81,7 @@ pub trait UpdateParams> { optimizer: &mut Optim, gradients: &Gradients, missing_tensors: &mut Vec, - ) -> Result<(), D::Err>; + ) -> Result<(), crate::tensor::Error>; } impl> UpdateParams for Tensor { @@ -118,7 +90,7 @@ impl> UpdateParams for Tensor { optimizer: &mut Optim, gradients: &Gradients, missing_tensors: &mut Vec, - ) -> Result<(), ::Err> { + ) -> Result<(), crate::tensor::Error> { optimizer.update_tensor(self, gradients, missing_tensors) } } @@ -128,12 +100,12 @@ pub trait ZeroGrads> { fn zero_grads(&self, grads: &mut Gradients) { self.try_zero_grads(grads).unwrap() } - fn try_zero_grads(&self, grads: &mut Gradients) -> Result<(), D::Err>; + fn try_zero_grads(&self, grads: &mut Gradients) -> Result<(), crate::tensor::Error>; fn alloc_grads(&self) -> Gradients { self.try_alloc_grads().unwrap() } - fn try_alloc_grads(&self) -> Result, D::Err> { + fn try_alloc_grads(&self) -> Result, crate::tensor::Error> { let mut grads = Gradients::leaky(); self.try_zero_grads(&mut grads)?; grads.retain_current_grads_as_leafs(); @@ -275,7 +247,7 @@ pub trait BuildModuleExt: Sized { self.try_build_module(m).unwrap() } - fn try_build_module(&self, m: M) -> Result + fn try_build_module(&self, m: M) -> Result where M: BuildOnDevice, M::Built: ResetParams, diff --git a/dfdx-core/src/nn_traits/tuples.rs b/dfdx-core/src/nn_traits/tuples.rs index 7cb587768..97e8c7deb 100644 --- a/dfdx-core/src/nn_traits/tuples.rs +++ b/dfdx-core/src/nn_traits/tuples.rs @@ -1,4 +1,8 @@ -use crate::{dtypes::Dtype, tensor::UniqueId, tensor_ops::Device}; +use crate::{ + dtypes::Dtype, + tensor::{Error, UniqueId}, + tensor_ops::Device, +}; use std::vec::Vec; @@ -7,7 +11,7 @@ macro_rules! tuple_impls { impl, Elem: Dtype, $($name: crate::nn_traits::BuildOnDevice),+> crate::nn_traits::BuildOnDevice for ($($name,)+) { type Built = ($($name::Built, )+); - fn try_build_on_device(&self, device: &Dev) -> Result { + fn try_build_on_device(&self, device: &Dev) -> Result { Ok(($( self.$idx.try_build_on_device(device)?, )+)) @@ -38,7 +42,7 @@ macro_rules! tuple_impls { } impl, Elem: Dtype, $($name: crate::nn_traits::ResetParams),+> crate::nn_traits::ResetParams for ($($name,)+) { - fn try_reset_params(&mut self) -> Result<(), Dev::Err> { + fn try_reset_params(&mut self) -> Result<(), Error> { $(self.$idx.try_reset_params()?;)+ Ok(()) } @@ -50,14 +54,14 @@ macro_rules! tuple_impls { optimizer: &mut Optim, gradients: &crate::prelude::Gradients, missing_tensors: &mut Vec, - ) -> Result<(), Dev::Err> { + ) -> Result<(), Error> { $(self.$idx.try_update_params(optimizer, gradients, missing_tensors)?;)+ Ok(()) } } impl, Elem: Dtype, $($name: crate::nn_traits::ZeroGrads),+> crate::nn_traits::ZeroGrads for ($($name,)+) { - fn try_zero_grads(&self, grads: &mut crate::prelude::Gradients) -> Result<(), Dev::Err> { + fn try_zero_grads(&self, grads: &mut crate::prelude::Gradients) -> Result<(), Error> { $(self.$idx.try_zero_grads(grads)?;)+ Ok(()) } @@ -91,20 +95,19 @@ macro_rules! tuple_impls { impl< Input, $last: - $(crate::nn_traits::Module::<$rev_tail ::Output, Error=$rev_tail::Error>, $rev_tail: )* + $(crate::nn_traits::Module::<$rev_tail ::Output>, $rev_tail: )* crate::nn_traits::Module > crate::nn_traits::Module for ($($name,)+) { type Output = $last ::Output; - type Error = $last ::Error; /// Calls forward sequentially on each module in the tuple. - fn try_forward(&self, x: Input) -> Result { + fn try_forward(&self, x: Input) -> Result { $(let x = self.$idx.try_forward(x)?;)+ Ok(x) } /// Calls forward sequentially on each module in the tuple. - fn try_forward_mut(&mut self, x: Input) -> Result { + fn try_forward_mut(&mut self, x: Input) -> Result { $(let x = self.$idx.try_forward_mut(x)?;)+ Ok(x) } diff --git a/dfdx-core/src/nn_traits/vecs.rs b/dfdx-core/src/nn_traits/vecs.rs index e413bc47c..803a07d8a 100644 --- a/dfdx-core/src/nn_traits/vecs.rs +++ b/dfdx-core/src/nn_traits/vecs.rs @@ -1,4 +1,8 @@ -use crate::{dtypes::Dtype, tensor::UniqueId, tensor_ops::Device}; +use crate::{ + dtypes::Dtype, + tensor::{Error, UniqueId}, + tensor_ops::Device, +}; use std::vec::Vec; @@ -6,7 +10,7 @@ impl, T: crate::nn_traits::BuildOnDevice> crate::nn_traits::BuildOnDevice for Vec { type Built = Vec; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { self.iter() .map(|m_i| m_i.try_build_on_device(device)) .collect() @@ -16,7 +20,7 @@ impl, T: crate::nn_traits::BuildOnDevice> impl, T: crate::nn_traits::ResetParams> crate::nn_traits::ResetParams for Vec { - fn try_reset_params(&mut self) -> Result<(), ::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { for m_i in self.iter_mut() { m_i.try_reset_params()?; } @@ -32,7 +36,7 @@ impl, T: crate::nn_traits::UpdateParams> optimizer: &mut Optim, gradients: &crate::tensor::Gradients, missing_tensors: &mut Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { for m_i in self.iter_mut() { m_i.try_update_params(optimizer, gradients, missing_tensors)?; } @@ -43,7 +47,10 @@ impl, T: crate::nn_traits::UpdateParams> impl, T: crate::nn_traits::ZeroGrads> crate::nn_traits::ZeroGrads for Vec { - fn try_zero_grads(&self, grads: &mut crate::tensor::Gradients) -> Result<(), ::Err> { + fn try_zero_grads( + &self, + grads: &mut crate::tensor::Gradients, + ) -> Result<(), crate::tensor::Error> { for m_i in self.iter() { m_i.try_zero_grads(grads)?; } @@ -82,15 +89,14 @@ impl> crate::nn_traits for Vec { type Output = T::Output; - type Error = T::Error; - fn try_forward(&self, mut x: Input) -> Result { + fn try_forward(&self, mut x: Input) -> Result { for m_i in self.iter() { x = m_i.try_forward(x)?; } Ok(x) } - fn try_forward_mut(&mut self, mut x: Input) -> Result { + fn try_forward_mut(&mut self, mut x: Input) -> Result { for m_i in self.iter_mut() { x = m_i.try_forward_mut(x)?; } diff --git a/dfdx-core/src/tensor/cpu/allocate.rs b/dfdx-core/src/tensor/cpu/allocate.rs index be134f127..cf93623b7 100644 --- a/dfdx-core/src/tensor/cpu/allocate.rs +++ b/dfdx-core/src/tensor/cpu/allocate.rs @@ -2,20 +2,17 @@ use crate::{ shapes::*, - tensor::{masks::triangle_mask, storage_traits::*, unique_id, Tensor}, + tensor::{masks::triangle_mask, storage_traits::*, unique_id, Error, Tensor}, }; -use super::{CachableVec, Cpu, CpuError, LendingIterator}; +use super::{CachableVec, Cpu, LendingIterator}; use rand::{distributions::Distribution, Rng}; use std::{sync::Arc, vec::Vec}; impl Cpu { #[inline] - pub(crate) fn try_alloc_zeros( - &self, - numel: usize, - ) -> Result, CpuError> { + pub(crate) fn try_alloc_zeros(&self, numel: usize) -> Result, Error> { self.try_alloc_elem::(numel, Default::default()) } @@ -24,14 +21,14 @@ impl Cpu { &self, numel: usize, elem: E, - ) -> Result, CpuError> { - let data = self.cache.try_pop::(numel).map_or_else( + ) -> Result, Error> { + let data: Result, Error> = self.cache.try_pop::(numel).map_or_else( #[cfg(feature = "fast-alloc")] || Ok(std::vec![elem; numel]), #[cfg(not(feature = "fast-alloc"))] || { let mut data: Vec = Vec::new(); - data.try_reserve(numel).map_err(|_| CpuError::OutOfMemory)?; + data.try_reserve(numel).map_err(|_| Error::OutOfMemory)?; data.resize(numel, elem); Ok(data) }, @@ -48,17 +45,17 @@ impl Cpu { data.fill(elem); Ok(data) }, - )?; + ); Ok(CachableVec { - data, + data: data?, cache: self.cache.clone(), }) } } impl ZerosTensor for Cpu { - fn try_zeros_like(&self, src: &S) -> Result, Self::Err> { + fn try_zeros_like(&self, src: &S) -> Result, Error> { let shape = *src.shape(); let strides = shape.strides(); let data = self.try_alloc_zeros::(shape.num_elements())?; @@ -75,14 +72,14 @@ impl ZerosTensor for Cpu { } impl ZeroFillStorage for Cpu { - fn try_fill_with_zeros(&self, storage: &mut Self::Vec) -> Result<(), Self::Err> { + fn try_fill_with_zeros(&self, storage: &mut Self::Vec) -> Result<(), Error> { storage.fill(Default::default()); Ok(()) } } impl OnesTensor for Cpu { - fn try_ones_like(&self, src: &S) -> Result, Self::Err> { + fn try_ones_like(&self, src: &S) -> Result, Error> { let shape = *src.shape(); let strides = shape.strides(); let data = self.try_alloc_elem::(shape.num_elements(), E::ONE)?; @@ -104,7 +101,7 @@ impl TriangleTensor for Cpu { src: &S, val: E, diagonal: impl Into>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let shape = *src.shape(); let strides = shape.strides(); let mut data = self.try_alloc_elem::(shape.num_elements(), val)?; @@ -126,7 +123,7 @@ impl TriangleTensor for Cpu { src: &S, val: E, diagonal: impl Into>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let shape = *src.shape(); let strides = shape.strides(); let mut data = self.try_alloc_elem::(shape.num_elements(), val)?; @@ -145,7 +142,7 @@ impl TriangleTensor for Cpu { } impl OneFillStorage for Cpu { - fn try_fill_with_ones(&self, storage: &mut Self::Vec) -> Result<(), Self::Err> { + fn try_fill_with_ones(&self, storage: &mut Self::Vec) -> Result<(), Error> { storage.fill(E::ONE); Ok(()) } @@ -156,7 +153,7 @@ impl SampleTensor for Cpu { &self, src: &S, distr: D, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut tensor = self.try_zeros_like(src)?; { #[cfg(not(feature = "no-std"))] @@ -173,7 +170,7 @@ impl SampleTensor for Cpu { &self, storage: &mut Self::Vec, distr: D, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { { #[cfg(not(feature = "no-std"))] let mut rng = self.rng.lock().unwrap(); @@ -201,11 +198,11 @@ impl TensorFromVec for Cpu { &self, src: Vec, shape: S, - ) -> Result, Self::Err> { + ) -> Result, Error> { let num_elements = shape.num_elements(); if src.len() != num_elements { - Err(CpuError::WrongNumElements) + Err(Error::WrongNumElements) } else { let src = CachableVec { data: src, diff --git a/dfdx-core/src/tensor/cpu/device.rs b/dfdx-core/src/tensor/cpu/device.rs index ebc380dcd..d3ce936f1 100644 --- a/dfdx-core/src/tensor/cpu/device.rs +++ b/dfdx-core/src/tensor/cpu/device.rs @@ -1,5 +1,5 @@ use crate::shapes::{Shape, Unit}; -use crate::tensor::{cache::TensorCache, cpu::LendingIterator, storage_traits::*, Tensor}; +use crate::tensor::{cache::TensorCache, cpu::LendingIterator, storage_traits::*, Error, Tensor}; use rand::{rngs::StdRng, Rng, SeedableRng}; use std::{sync::Arc, vec::Vec}; @@ -47,30 +47,6 @@ impl Cpu { } } -#[derive(Debug, Clone, Copy)] -pub enum CpuError { - /// Device is out of memory - OutOfMemory, - /// Not enough elements were provided when creating a tensor - WrongNumElements, -} - -impl std::fmt::Display for CpuError { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - match self { - Self::OutOfMemory => f.write_str("CpuError::OutOfMemory"), - Self::WrongNumElements => f.write_str("CpuError::WrongNumElements"), - } - } -} - -#[cfg(feature = "std")] -impl std::error::Error for CpuError {} - -impl HasErr for Cpu { - type Err = CpuError; -} - /// A [Vec] that can be cloned without allocating new memory. /// When [Drop]ed it will insert it's data into the cache. #[derive(Debug)] @@ -154,7 +130,7 @@ impl RandomU64 for Cpu { impl Storage for Cpu { type Vec = CachableVec; - fn try_alloc_len(&self, len: usize) -> Result { + fn try_alloc_len(&self, len: usize) -> Result { self.try_alloc_zeros(len) } @@ -173,23 +149,23 @@ impl Storage for Cpu { } impl Synchronize for Cpu { - fn try_synchronize(&self) -> Result<(), Self::Err> { + fn try_synchronize(&self) -> Result<(), Error> { Ok(()) } } impl Cache for Cpu { - fn try_enable_cache(&self) -> Result<(), Self::Err> { + fn try_enable_cache(&self) -> Result<(), Error> { self.cache.enable(); Ok(()) } - fn try_disable_cache(&self) -> Result<(), Self::Err> { + fn try_disable_cache(&self) -> Result<(), Error> { self.cache.disable(); self.try_empty_cache() } - fn try_empty_cache(&self) -> Result<(), Self::Err> { + fn try_empty_cache(&self) -> Result<(), Error> { #[cfg(not(feature = "no-std"))] let mut cache = self.cache.allocations.write().unwrap(); #[cfg(feature = "no-std")] diff --git a/dfdx-core/src/tensor/cpu/mod.rs b/dfdx-core/src/tensor/cpu/mod.rs index fda5306ea..f69ef5dd0 100644 --- a/dfdx-core/src/tensor/cpu/mod.rs +++ b/dfdx-core/src/tensor/cpu/mod.rs @@ -7,7 +7,7 @@ pub(crate) use index::index_to_i; pub(crate) use iterate::{LendingIterator, NdIndex}; pub(crate) use device::CachableVec; -pub use device::{Cpu, CpuError}; +pub use device::Cpu; #[cfg(test)] mod tests { diff --git a/dfdx-core/src/tensor/cuda/allocate.rs b/dfdx-core/src/tensor/cuda/allocate.rs index fb688ae37..aa489f9ad 100644 --- a/dfdx-core/src/tensor/cuda/allocate.rs +++ b/dfdx-core/src/tensor/cuda/allocate.rs @@ -2,10 +2,10 @@ use crate::{ shapes::*, - tensor::{masks::triangle_mask, storage_traits::*, unique_id, Cpu, CpuError, NoneTape, Tensor}, + tensor::{masks::triangle_mask, storage_traits::*, unique_id, Cpu, Error, NoneTape, Tensor}, }; -use super::{device::CachableCudaSlice, Cuda, CudaError}; +use super::{device::CachableCudaSlice, Cuda}; use cudarc::driver::{CudaSlice, DeviceSlice}; use rand::Rng; @@ -16,7 +16,7 @@ impl Cuda { &self, shape: S, buf: Vec, - ) -> Result, CudaError> { + ) -> Result, Error> { let mut slice = unsafe { self.alloc_empty(buf.len()) }?; self.dev.htod_copy_into(buf, &mut slice)?; Ok(self.build_tensor(shape, shape.strides(), slice)) @@ -44,7 +44,7 @@ impl Cuda { } impl ZerosTensor for Cuda { - fn try_zeros_like(&self, src: &S) -> Result, Self::Err> { + fn try_zeros_like(&self, src: &S) -> Result, Error> { let shape = *src.shape(); let strides = shape.strides(); let mut data = unsafe { self.alloc_empty(shape.num_elements()) }?; @@ -54,7 +54,7 @@ impl ZerosTensor for Cuda { } impl ZeroFillStorage for Cuda { - fn try_fill_with_zeros(&self, storage: &mut Self::Vec) -> Result<(), Self::Err> { + fn try_fill_with_zeros(&self, storage: &mut Self::Vec) -> Result<(), Error> { self.dev.memset_zeros(&mut storage.data)?; Ok(()) } @@ -64,7 +64,7 @@ impl OnesTensor for Cuda where Cpu: OnesTensor, { - fn try_ones_like(&self, src: &S) -> Result, Self::Err> { + fn try_ones_like(&self, src: &S) -> Result, Error> { let shape = *src.shape(); let buf = std::vec![E::ONE; shape.num_elements()]; self.tensor_from_host_buf(shape, buf) @@ -80,7 +80,7 @@ where src: &S, val: E, diagonal: impl Into>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let shape = *src.shape(); let mut data = std::vec![val; shape.num_elements()]; let offset = diagonal.into().unwrap_or(0); @@ -93,7 +93,7 @@ where src: &S, val: E, diagonal: impl Into>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let shape = *src.shape(); let mut data = std::vec![val; shape.num_elements()]; let offset = diagonal.into().unwrap_or(0); @@ -103,7 +103,7 @@ where } impl OneFillStorage for Cuda { - fn try_fill_with_ones(&self, storage: &mut Self::Vec) -> Result<(), Self::Err> { + fn try_fill_with_ones(&self, storage: &mut Self::Vec) -> Result<(), Error> { self.dev .htod_copy_into(std::vec![E::ONE; storage.len()], &mut storage.data)?; Ok(()) @@ -118,7 +118,7 @@ where &self, src: &S, distr: D, - ) -> Result, Self::Err> { + ) -> Result, Error> { let shape = *src.shape(); let mut buf = Vec::with_capacity(shape.num_elements()); { @@ -134,7 +134,7 @@ where &self, storage: &mut Self::Vec, distr: D, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut buf = Vec::with_capacity(storage.len()); { #[cfg(not(feature = "no-std"))] @@ -180,11 +180,11 @@ impl TensorFromVec for Cuda { &self, src: Vec, shape: S, - ) -> Result, Self::Err> { + ) -> Result, Error> { let num_elements = shape.num_elements(); if src.len() != num_elements { - Err(CudaError::Cpu(CpuError::WrongNumElements)) + Err(Error::WrongNumElements) } else { self.tensor_from_host_buf(shape, src) } diff --git a/dfdx-core/src/tensor/cuda/device.rs b/dfdx-core/src/tensor/cuda/device.rs index 23a2887a4..156bdddd7 100644 --- a/dfdx-core/src/tensor/cuda/device.rs +++ b/dfdx-core/src/tensor/cuda/device.rs @@ -1,7 +1,7 @@ use crate::shapes::{Shape, Unit}; -use crate::tensor::cpu::{Cpu, CpuError}; +use crate::tensor::cpu::Cpu; use crate::tensor::{ - cache::TensorCache, Cache, HasErr, NoneTape, RandomU64, Storage, Synchronize, Tensor, + cache::TensorCache, Cache, Error, NoneTape, RandomU64, Storage, Synchronize, Tensor, }; use cudarc::driver::{DevicePtr, DevicePtrMut, DeviceRepr}; @@ -32,37 +32,22 @@ pub struct Cuda { pub(crate) cache: Arc>, } -#[derive(Debug)] -pub enum CudaError { - Blas(CublasError), - #[cfg(feature = "cudnn")] - Cudnn(cudarc::cudnn::CudnnError), - Driver(DriverError), - Cpu(CpuError), -} - -impl From for CudaError { - fn from(value: CpuError) -> Self { - Self::Cpu(value) - } -} - -impl From for CudaError { +impl From for Error { fn from(value: CublasError) -> Self { - Self::Blas(value) + Self::CublasError(value) } } -impl From for CudaError { +impl From for Error { fn from(value: DriverError) -> Self { - Self::Driver(value) + Self::CudaDriverError(value) } } #[cfg(feature = "cudnn")] -impl From for CudaError { +impl From for Error { fn from(value: cudarc::cudnn::CudnnError) -> Self { - Self::Cudnn(value) + Self::CudnnError(value) } } @@ -79,12 +64,12 @@ impl Cuda { } /// Constructs rng with the given seed. - pub fn try_seed_from_u64(seed: u64) -> Result { + pub fn try_seed_from_u64(seed: u64) -> Result { Self::try_build(0, seed) } /// Constructs with the given seed & device ordinal - pub fn try_build(ordinal: usize, seed: u64) -> Result { + pub fn try_build(ordinal: usize, seed: u64) -> Result { let cpu = Cpu::seed_from_u64(seed); let dev = CudaDevice::new(ordinal)?; let blas = Arc::new(CudaBlas::new(dev.clone())?); @@ -112,7 +97,7 @@ impl Cuda { pub(crate) unsafe fn alloc_empty( &self, len: usize, - ) -> Result, CudaError> { + ) -> Result, Error> { let data = self.cache.try_pop::(len).map_or_else( || self.dev.alloc::(len), |ptr| Ok(self.dev.upgrade_device_ptr(ptr, len)), @@ -123,7 +108,7 @@ impl Cuda { pub(crate) unsafe fn get_workspace( &self, len: usize, - ) -> Result>, CudaError> { + ) -> Result>, Error> { let num_bytes_required = len * std::mem::size_of::(); let mut workspace = self.workspace.as_ref().lock().unwrap(); @@ -137,16 +122,6 @@ impl Cuda { } } -impl std::fmt::Display for CudaError { - fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { - write!(f, "{self:?}") - } -} - -impl HasErr for Cuda { - type Err = CudaError; -} - /// A [CudaSlice] that can be cloned without allocating new memory. /// When [Drop]ed it will insert it's data into the cache. #[derive(Debug)] @@ -246,17 +221,17 @@ impl RandomU64 for Cuda { } impl Cache for Cuda { - fn try_enable_cache(&self) -> Result<(), Self::Err> { + fn try_enable_cache(&self) -> Result<(), Error> { self.cache.enable(); Ok(()) } - fn try_disable_cache(&self) -> Result<(), Self::Err> { + fn try_disable_cache(&self) -> Result<(), Error> { self.cache.disable(); self.try_empty_cache() } - fn try_empty_cache(&self) -> Result<(), Self::Err> { + fn try_empty_cache(&self) -> Result<(), Error> { #[cfg(not(feature = "no-std"))] let mut cache = self.cache.allocations.write().unwrap(); #[cfg(feature = "no-std")] @@ -273,15 +248,15 @@ impl Cache for Cuda { } impl Synchronize for Cuda { - fn try_synchronize(&self) -> Result<(), CudaError> { - self.dev.synchronize().map_err(CudaError::from) + fn try_synchronize(&self) -> Result<(), Error> { + self.dev.synchronize().map_err(Error::from) } } impl Storage for Cuda { type Vec = CachableCudaSlice; - fn try_alloc_len(&self, len: usize) -> Result { + fn try_alloc_len(&self, len: usize) -> Result { let mut data = unsafe { self.alloc_empty(len) }?; self.dev.memset_zeros(&mut data)?; Ok(CachableCudaSlice { diff --git a/dfdx-core/src/tensor/cuda/mod.rs b/dfdx-core/src/tensor/cuda/mod.rs index 8b91d2ab3..a38149cfc 100644 --- a/dfdx-core/src/tensor/cuda/mod.rs +++ b/dfdx-core/src/tensor/cuda/mod.rs @@ -1,7 +1,7 @@ mod allocate; mod device; -pub use device::{Cuda, CudaError}; +pub use device::Cuda; pub(crate) fn launch_cfg(n: u32) -> cudarc::driver::LaunchConfig { let num_blocks = (n + NUM_THREADS - 1) / NUM_THREADS; diff --git a/dfdx-core/src/tensor/error.rs b/dfdx-core/src/tensor/error.rs new file mode 100644 index 000000000..f6b43c322 --- /dev/null +++ b/dfdx-core/src/tensor/error.rs @@ -0,0 +1,28 @@ +/// Represents a number of different errors that can occur from creating tensors +/// or launching tensor operations. This encompasses both Cpu and CUDA errors. +#[non_exhaustive] +#[derive(Debug)] +pub enum Error { + /// Device is out of memory + OutOfMemory, + /// Not enough elements were provided when creating a tensor + WrongNumElements, + /// Some tensors were unused by an optimizer in a graph. + UnusedTensors(std::vec::Vec), + #[cfg(feature = "cuda")] + CublasError(cudarc::cublas::result::CublasError), + #[cfg(feature = "cuda")] + CudaDriverError(cudarc::driver::DriverError), + + #[cfg(feature = "cudnn")] + CudnnError(cudarc::cudnn::CudnnError), +} + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + write!(f, "{self:?}") + } +} + +#[cfg(feature = "std")] +impl std::error::Error for Error {} diff --git a/dfdx-core/src/tensor/ghost.rs b/dfdx-core/src/tensor/ghost.rs index c05854ab6..8c349a8fe 100644 --- a/dfdx-core/src/tensor/ghost.rs +++ b/dfdx-core/src/tensor/ghost.rs @@ -42,10 +42,6 @@ impl> Clone for GhostTensor { } } -impl> super::storage_traits::HasErr for GhostTensor { - type Err = D::Err; -} - impl> HasShape for GhostTensor { type WithShape = GhostTensor; type Shape = S; @@ -56,7 +52,7 @@ impl> HasShape for GhostTensor { impl> super::storage_traits::AllocGrad for GhostTensor { type Gradient = D::Vec; - fn try_alloc_grad(&self) -> Result { + fn try_alloc_grad(&self) -> Result { self.dev.try_alloc_len(self.len) } } diff --git a/dfdx-core/src/tensor/gradients.rs b/dfdx-core/src/tensor/gradients.rs index 99dc7a163..86974ec61 100644 --- a/dfdx-core/src/tensor/gradients.rs +++ b/dfdx-core/src/tensor/gradients.rs @@ -5,7 +5,7 @@ use std::collections::{BTreeMap, BTreeSet}; use std::{boxed::Box, vec::Vec}; use super::tensorlike::Tensorlike; -use super::{storage_traits::Storage, unique_id, Tensor, UniqueId}; +use super::{storage_traits::Storage, unique_id, Error, Tensor, UniqueId}; use crate::shapes::Shape; /// A generic container for keeping gradients of tensors keyed by the @@ -44,13 +44,13 @@ impl> Gradients { pub fn get_or_alloc_mut( &mut self, t: &impl Tensorlike, - ) -> Result<&mut D::Vec, D::Err> { + ) -> Result<&mut D::Vec, Error> { self.try_alloc_for(t)?; Ok(self.get_mut(t)) } /// Inserts a gradient for `t` - pub fn try_alloc_for(&mut self, t: &impl Tensorlike) -> Result<(), D::Err> { + pub fn try_alloc_for(&mut self, t: &impl Tensorlike) -> Result<(), Error> { if let std::collections::btree_map::Entry::Vacant(e) = self.gradient_by_id.entry(t.id()) { e.insert(t.try_alloc_grad()?); } @@ -179,7 +179,7 @@ impl> Gradients { pub struct OwnedTape> { /// A list of (Time, BackwardOp) pairs. The Time is used to ensure operations /// from merged tapes are executed in the correct order. - pub(crate) operations: Vec<(UniqueId, BackwardOp)>, + pub(crate) operations: Vec<(UniqueId, BackwardOp)>, pub(crate) gradients: Gradients, } @@ -214,7 +214,7 @@ impl> OwnedTape { /// Compute the [Gradients]! This just runs all the operations on a new [Gradients] struct. /// /// Note that this method takes ownership of self, so it can't be called twice! - pub(crate) fn execute(&mut self) -> Result, D::Err> { + pub(crate) fn execute(&mut self) -> Result, Error> { // We must ensure that the operations are sorted in execution time order. // Otherwise an backward operation may not be executed in the right order // if multiple tapes were merged together. @@ -228,7 +228,7 @@ impl> OwnedTape { } } -type BackwardOp = Box) -> Result<(), Err>>; +type BackwardOp = Box) -> Result<(), Error>>; /// Contains nothing. When [Tape::add_backward_op] is called, this struct does nothing. #[derive(Default, Debug, Clone, Copy)] @@ -240,14 +240,14 @@ pub trait Tape>: Default + Merge + Merge { const OWNS_TAPE: bool; fn add_backward_op(&mut self, operation: F) where - F: 'static + FnOnce(&mut Gradients) -> Result<(), D::Err>; + F: 'static + FnOnce(&mut Gradients) -> Result<(), Error>; } impl> Tape for OwnedTape { const OWNS_TAPE: bool = true; fn add_backward_op(&mut self, operation: F) where - F: 'static + FnOnce(&mut Gradients) -> Result<(), D::Err>, + F: 'static + FnOnce(&mut Gradients) -> Result<(), Error>, { self.operations.push((unique_id(), Box::new(operation))); } @@ -257,7 +257,7 @@ impl> Tape for NoneTape { const OWNS_TAPE: bool = false; fn add_backward_op(&mut self, _: F) where - F: 'static + FnOnce(&mut Gradients) -> Result<(), D::Err>, + F: 'static + FnOnce(&mut Gradients) -> Result<(), Error>, { } } @@ -329,7 +329,7 @@ impl> Tape for std::sync::Arc(&mut self, operation: F) where - F: 'static + FnOnce(&mut Gradients) -> Result<(), D::Err>, + F: 'static + FnOnce(&mut Gradients) -> Result<(), Error>, { let mut tape = self.lock().unwrap(); tape.add_backward_op(operation); diff --git a/dfdx-core/src/tensor/mod.rs b/dfdx-core/src/tensor/mod.rs index 1b3a5b06d..2a7f9db3a 100644 --- a/dfdx-core/src/tensor/mod.rs +++ b/dfdx-core/src/tensor/mod.rs @@ -147,6 +147,7 @@ mod masks; pub(crate) mod numpy; #[cfg(feature = "numpy")] pub use numpy::NumpyDtype; +mod error; #[cfg(feature = "safetensors")] pub mod safetensors; mod tensorlike; @@ -155,23 +156,24 @@ mod unique_id; pub(crate) mod storage_traits; mod tensor_impls; +pub use error::Error; pub(crate) use ghost::GhostTensor; pub(crate) use storage_traits::{OneFillStorage, ZeroFillStorage}; pub use tensorlike::Tensorlike; -pub use cpu::{Cpu, CpuError}; +pub use cpu::Cpu; #[cfg(not(feature = "cuda"))] pub type AutoDevice = Cpu; #[cfg(feature = "cuda")] pub(crate) use cuda::launch_cfg; #[cfg(feature = "cuda")] -pub use cuda::{Cuda, CudaError}; +pub use cuda::Cuda; #[cfg(feature = "cuda")] pub type AutoDevice = Cuda; pub use storage_traits::{AsArray, CopySlice, TensorFrom, TensorFromVec, TensorToArray}; -pub use storage_traits::{Cache, HasErr, RandomU64, Storage, Synchronize}; +pub use storage_traits::{Cache, RandomU64, Storage, Synchronize}; pub use storage_traits::{OnesTensor, SampleTensor, TriangleTensor, ZerosTensor}; pub use tensor_impls::{PutTape, SplitTape, Tensor, Trace, WithEmptyTape}; diff --git a/dfdx-core/src/tensor/storage_traits.rs b/dfdx-core/src/tensor/storage_traits.rs index 4f98c3e41..9578947d9 100644 --- a/dfdx-core/src/tensor/storage_traits.rs +++ b/dfdx-core/src/tensor/storage_traits.rs @@ -4,12 +4,7 @@ use std::vec::Vec; use crate::shapes::*; -use super::Tensor; - -/// Represents something that has an error associated type -pub trait HasErr: Sized { - type Err: std::fmt::Debug + std::fmt::Display; -} +use super::{Error, Tensor}; pub trait RandomU64 { /// Generates a random u64 number @@ -17,40 +12,40 @@ pub trait RandomU64 { } /// Something that can store nd arrays for a given [Shape] and [Dtype] -pub trait Storage: 'static + std::fmt::Debug + Default + Clone + HasErr { +pub trait Storage: 'static + std::fmt::Debug + Default + Clone { /// Generic Storage type type Vec: 'static + std::fmt::Debug + Clone + Send + Sync; /// Allocates a gradient for the given nd array - fn try_alloc_grad(&self, storage: &Self::Vec) -> Result { + fn try_alloc_grad(&self, storage: &Self::Vec) -> Result { self.try_alloc_len(self.len(storage)) } - fn try_alloc_len(&self, len: usize) -> Result; + fn try_alloc_len(&self, len: usize) -> Result; fn tensor_to_vec(&self, tensor: &Tensor) -> Vec; fn len(&self, v: &Self::Vec) -> usize; } -pub trait Synchronize: HasErr { +pub trait Synchronize { /// Blocks until all work on device to complete. Useful for benchmarking. fn synchronize(&self) { self.try_synchronize().unwrap() } /// Blocks until all work on device to complete. Useful for benchmarking. - fn try_synchronize(&self) -> Result<(), Self::Err>; + fn try_synchronize(&self) -> Result<(), Error>; } -pub trait Cache: HasErr { +pub trait Cache { /// Enables the cache of the device. fn enable_cache(&self) { self.try_enable_cache().unwrap() } /// Tries to enable the cache of the device. - fn try_enable_cache(&self) -> Result<(), Self::Err>; + fn try_enable_cache(&self) -> Result<(), Error>; /// Disables the cache of the device. This will also empty the cache /// if there are things in it. See [Cache::empty_cache] for @@ -61,7 +56,7 @@ pub trait Cache: HasErr { /// Tries to disable the cache of the device. See [Cache::disable_cache] for /// details of when this is useful. - fn try_disable_cache(&self) -> Result<(), Self::Err>; + fn try_disable_cache(&self) -> Result<(), Error>; /// Empties the cache of the device. /// @@ -78,18 +73,18 @@ pub trait Cache: HasErr { /// Tries to empty the cache of the device. See [Cache::empty_cache] for /// details of when this is useful. - fn try_empty_cache(&self) -> Result<(), Self::Err>; + fn try_empty_cache(&self) -> Result<(), Error>; } /// Internal trait - Represents something that can allocate its own gradient. -pub trait AllocGrad: HasErr { +pub trait AllocGrad { type Gradient: 'static; - fn try_alloc_grad(&self) -> Result; + fn try_alloc_grad(&self) -> Result; } impl, T> AllocGrad for Tensor { type Gradient = D::Vec; - fn try_alloc_grad(&self) -> Result { + fn try_alloc_grad(&self) -> Result { self.device.try_alloc_grad(self.data.as_ref()) } } @@ -143,7 +138,7 @@ pub trait ZerosTensor: Storage { } /// Fallible version of [ZerosTensor::zeros] - fn try_zeros(&self) -> Result, Self::Err> { + fn try_zeros(&self) -> Result, Error> { self.try_zeros_like::(&Default::default()) } @@ -168,11 +163,11 @@ pub trait ZerosTensor: Storage { } /// Fallible version of [ZerosTensor::zeros_like] - fn try_zeros_like(&self, src: &S) -> Result, Self::Err>; + fn try_zeros_like(&self, src: &S) -> Result, Error>; } pub trait ZeroFillStorage: Storage { - fn try_fill_with_zeros(&self, storage: &mut Self::Vec) -> Result<(), Self::Err>; + fn try_fill_with_zeros(&self, storage: &mut Self::Vec) -> Result<(), Error>; } /// Construct tensors filled with ones. @@ -188,7 +183,7 @@ pub trait OnesTensor: Storage { } /// Fallible version of [OnesTensor::ones] - fn try_ones(&self) -> Result, Self::Err> { + fn try_ones(&self) -> Result, Error> { self.try_ones_like::(&Default::default()) } @@ -213,11 +208,11 @@ pub trait OnesTensor: Storage { } /// Fallible version of [OnesTensor::ones_like] - fn try_ones_like(&self, src: &S) -> Result, Self::Err>; + fn try_ones_like(&self, src: &S) -> Result, Error>; } pub trait OneFillStorage: Storage { - fn try_fill_with_ones(&self, storage: &mut Self::Vec) -> Result<(), Self::Err>; + fn try_fill_with_ones(&self, storage: &mut Self::Vec) -> Result<(), Error>; } /// Build upper & lower triangle tensors. @@ -266,7 +261,7 @@ pub trait TriangleTensor: Storage { &self, val: E, diagonal: impl Into>, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.try_upper_tri_like::(&Default::default(), val, diagonal) } @@ -286,7 +281,7 @@ pub trait TriangleTensor: Storage { src: &S, val: E, diagonal: impl Into>, - ) -> Result, Self::Err>; + ) -> Result, Error>; /// Build a tensor containing the lower triangle part of each lowest 2D matrix /// set to the given value, along the given diagonal. The other values will be `E::default()`. @@ -332,7 +327,7 @@ pub trait TriangleTensor: Storage { &self, val: E, diagonal: impl Into>, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.try_lower_tri_like::(&Default::default(), val, diagonal) } @@ -352,7 +347,7 @@ pub trait TriangleTensor: Storage { src: &S, val: E, diagonal: impl Into>, - ) -> Result, Self::Err>; + ) -> Result, Error>; } /// Constructs tensors filled with random values from a given distribution. @@ -396,7 +391,7 @@ pub trait SampleTensor: Storage { fn try_sample>( &self, distr: D, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.try_sample_like::(&Default::default(), distr) } @@ -413,14 +408,14 @@ pub trait SampleTensor: Storage { &self, src: &S, distr: D, - ) -> Result, Self::Err>; + ) -> Result, Error>; /// Fills tensor `Storage` with data from a given distribution fn try_fill_with_distr>( &self, storage: &mut Self::Vec, distr: D, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } pub trait TensorToArray: Storage { @@ -457,7 +452,7 @@ pub trait TensorFromVec: Storage { &self, src: Vec, shape: S, - ) -> Result, Self::Err>; + ) -> Result, Error>; } impl, T> Tensor { @@ -470,7 +465,7 @@ impl, T> Tensor { pub fn try_to_device>( &self, device: &Dst, - ) -> Result, Dst::Err> { + ) -> Result, Error> { let buf = self.as_vec(); device.try_tensor_from_vec(buf, self.shape) } @@ -492,23 +487,23 @@ pub trait TensorFrom: Storage { self.try_tensor(src).unwrap() } /// Fallible version of [TensorFrom::tensor] - fn try_tensor(&self, src: Src) -> Result, Self::Err>; + fn try_tensor(&self, src: Src) -> Result, Error>; } impl> TensorFrom for D { - fn try_tensor(&self, src: E) -> Result, Self::Err> { + fn try_tensor(&self, src: E) -> Result, Error> { self.try_tensor_from_vec(vec![src], ()) } } impl> TensorFrom<[E; M], Rank1, E> for D { - fn try_tensor(&self, src: [E; M]) -> Result, E, Self>, Self::Err> { + fn try_tensor(&self, src: [E; M]) -> Result, E, Self>, Error> { self.try_tensor(&src) } } impl> TensorFrom<&[E; M], Rank1, E> for D { - fn try_tensor(&self, src: &[E; M]) -> Result, E, Self>, Self::Err> { + fn try_tensor(&self, src: &[E; M]) -> Result, E, Self>, Error> { self.try_tensor_from_vec(src.to_vec(), (Const::,)) } } @@ -516,7 +511,7 @@ impl> TensorFrom<&[E; M], Rank1, impl> TensorFrom<[[E; N]; M], Rank2, E> for D { - fn try_tensor(&self, src: [[E; N]; M]) -> Result, E, Self>, Self::Err> { + fn try_tensor(&self, src: [[E; N]; M]) -> Result, E, Self>, Error> { let vec: Vec = src.iter().flat_map(|v| v.iter().copied()).collect(); self.try_tensor_from_vec(vec, (Const::, Const::)) @@ -526,10 +521,7 @@ impl> impl> TensorFrom<[[[E; O]; N]; M], Rank3, E> for D { - fn try_tensor( - &self, - src: [[[E; O]; N]; M], - ) -> Result, E, Self>, Self::Err> { + fn try_tensor(&self, src: [[[E; O]; N]; M]) -> Result, E, Self>, Error> { let vec: Vec = src .iter() .flat_map(|v| v.iter()) @@ -552,7 +544,7 @@ impl< fn try_tensor( &self, src: [[[[E; P]; O]; N]; M], - ) -> Result, E, Self>, Self::Err> { + ) -> Result, E, Self>, Error> { let vec: Vec = src .iter() .flat_map(|v| v.iter()) @@ -565,13 +557,13 @@ impl< } impl> TensorFrom, S, E> for D { - fn try_tensor(&self, src: Vec) -> Result, Self::Err> { + fn try_tensor(&self, src: Vec) -> Result, Error> { self.try_tensor_from_vec(src, S::default()) } } impl> TensorFrom<(Vec, S), S, E> for D { - fn try_tensor(&self, (src, shape): (Vec, S)) -> Result, Self::Err> { + fn try_tensor(&self, (src, shape): (Vec, S)) -> Result, Error> { self.try_tensor_from_vec(src, shape) } } diff --git a/dfdx-core/src/tensor/tensor_impls.rs b/dfdx-core/src/tensor/tensor_impls.rs index bf4c57f22..691c5931f 100644 --- a/dfdx-core/src/tensor/tensor_impls.rs +++ b/dfdx-core/src/tensor/tensor_impls.rs @@ -54,10 +54,6 @@ impl, T> HasDtype for Tensor { type Dtype = E; } -impl, T> HasErr for Tensor { - type Err = D::Err; -} - /// Something that can trace gradients pub trait Trace>: Clone { type Traced; @@ -198,7 +194,7 @@ impl, T> Tensor { self.try_fill_with_zeros().unwrap() } /// Fallible version of [Tensor::fill_with_zeros] - pub fn try_fill_with_zeros(&mut self) -> Result<(), D::Err> { + pub fn try_fill_with_zeros(&mut self) -> Result<(), Error> { self.device .try_fill_with_zeros(Arc::make_mut(&mut self.data)) } @@ -210,7 +206,7 @@ impl, T> Tensor { self.try_fill_with_ones().unwrap() } /// Fallible version of [Tensor::fill_with_ones] - pub fn try_fill_with_ones(&mut self) -> Result<(), D::Err> { + pub fn try_fill_with_ones(&mut self) -> Result<(), Error> { self.device .try_fill_with_ones(Arc::make_mut(&mut self.data)) } @@ -226,7 +222,7 @@ impl, T> Tensor { pub fn try_fill_with_distr>( &mut self, distr: Distr, - ) -> Result<(), D::Err> { + ) -> Result<(), Error> { self.device .try_fill_with_distr(Arc::make_mut(&mut self.data), distr) } diff --git a/dfdx-core/src/tensor/tensorlike.rs b/dfdx-core/src/tensor/tensorlike.rs index 18e4f074d..5b33f9bd5 100644 --- a/dfdx-core/src/tensor/tensorlike.rs +++ b/dfdx-core/src/tensor/tensorlike.rs @@ -1,8 +1,4 @@ -use crate::{ - prelude::{HasErr, HasShape}, - shapes::Shape, - tensor::Storage, -}; +use crate::{prelude::HasShape, shapes::Shape, tensor::Storage}; use super::{storage_traits::AllocGrad, GhostTensor, Tensor, UniqueId}; @@ -12,7 +8,7 @@ use super::{storage_traits::AllocGrad, GhostTensor, Tensor, UniqueId}; /// *If it looks like a tensor and barks like a tensor, then pet it like a tensor.* #[allow(clippy::len_without_is_empty)] pub trait Tensorlike>: - AllocGrad + HasErr + HasShape + AllocGrad + HasShape { fn id(&self) -> UniqueId; fn len(&self) -> usize; diff --git a/dfdx-core/src/tensor_ops/abs/mod.rs b/dfdx-core/src/tensor_ops/abs/mod.rs index 361079f34..f7ac117a4 100644 --- a/dfdx-core/src/tensor_ops/abs/mod.rs +++ b/dfdx-core/src/tensor_ops/abs/mod.rs @@ -34,7 +34,7 @@ impl, T: Tape> Tensor Result { + pub fn try_abs(self) -> Result { try_unary_op(AbsKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs b/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs index 56dda8a04..f409954ed 100644 --- a/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs +++ b/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs @@ -47,7 +47,7 @@ impl, T: Tape> self.try_accurate_gelu().unwrap() } /// See [accurate_gelu] - pub fn try_accurate_gelu(self) -> Result { + pub fn try_accurate_gelu(self) -> Result { try_unary_op(AccurateGeLUKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/adam/cpu_kernel.rs b/dfdx-core/src/tensor_ops/adam/cpu_kernel.rs index b89c0f4d2..998a578e1 100644 --- a/dfdx-core/src/tensor_ops/adam/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/adam/cpu_kernel.rs @@ -1,7 +1,7 @@ use super::{AdamConfig, AdamKernel, WeightDecay}; use crate::{ dtypes::{Dtype, NotMixedPrecision}, - tensor::Cpu, + tensor::{Cpu, Error}, }; #[cfg(feature = "f16")] @@ -14,7 +14,7 @@ impl AdamKernel> for Cpu { moment1: &mut Self::Vec, moment2: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let betas = cfg.betas.map(|x| x as f32); let eps = cfg.eps as f32; let lr = cfg.lr as f32; @@ -60,7 +60,7 @@ impl AdamKernel for Cpu { moment1: &mut Self::Vec, moment2: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let betas = cfg.betas.map(E::from_f64).map(Option::unwrap); let eps = E::from_f64(cfg.eps).unwrap(); let lr = E::from_f64(cfg.lr).unwrap(); diff --git a/dfdx-core/src/tensor_ops/adam/cuda_kernel.rs b/dfdx-core/src/tensor_ops/adam/cuda_kernel.rs index 3d1bfb9c7..24617b7c8 100644 --- a/dfdx-core/src/tensor_ops/adam/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/adam/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ dtypes::*, - tensor::{launch_cfg, Cuda}, + tensor::{launch_cfg, Cuda, Error}, tensor_ops::optim::*, }; @@ -72,7 +72,7 @@ where moment1: &mut Self::Vec, moment2: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FWD) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, &[Self::FWD])?; } diff --git a/dfdx-core/src/tensor_ops/adam/mod.rs b/dfdx-core/src/tensor_ops/adam/mod.rs index fd81e64cc..b9a307232 100644 --- a/dfdx-core/src/tensor_ops/adam/mod.rs +++ b/dfdx-core/src/tensor_ops/adam/mod.rs @@ -5,7 +5,7 @@ mod cuda_kernel; use crate::{ shapes::{Dtype, Shape}, - tensor::{Storage, Tensor}, + tensor::{Error, Storage, Tensor}, }; use super::WeightDecay; @@ -57,7 +57,7 @@ pub trait AdamKernel: Storage { moment1: &mut Self::Vec, moment2: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl AdamConfig { @@ -68,7 +68,7 @@ impl AdamConfig { moment1: &mut D::Vec, moment2: &mut D::Vec, grad: &D::Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { param.device.adam_kernel( t, self, diff --git a/dfdx-core/src/tensor_ops/add/mod.rs b/dfdx-core/src/tensor_ops/add/mod.rs index a70270042..29010e6f1 100644 --- a/dfdx-core/src/tensor_ops/add/mod.rs +++ b/dfdx-core/src/tensor_ops/add/mod.rs @@ -6,7 +6,7 @@ mod cuda_kernel; use super::ops::*; use crate::{ shapes::*, - tensor::{HasErr, Merge, Storage, Tape, Tensor}, + tensor::{Error, Merge, Storage, Tape, Tensor}, }; #[repr(C)] @@ -49,9 +49,9 @@ where } /// Fallible version of [std::ops::Add]. See [add] -pub trait TryAdd: HasErr { +pub trait TryAdd { type Output; - fn try_add(self, rhs: Rhs) -> Result; + fn try_add(self, rhs: Rhs) -> Result; } impl, R> TryAdd> @@ -62,7 +62,7 @@ where { type Output = Self; /// See [add] - fn try_add(self, rhs: Tensor) -> Result { + fn try_add(self, rhs: Tensor) -> Result { try_binary_op(BinaryAddKernelOp, self, rhs) } } @@ -73,7 +73,7 @@ where { type Output = Self; /// See [add] - fn try_add(self, rhs: Rhs) -> Result { + fn try_add(self, rhs: Rhs) -> Result { let rhs: f64 = rhs.into(); let scalar = E::from_f64(rhs).unwrap(); try_unary_op(ScalarAddKernelOp { scalar }, self) diff --git a/dfdx-core/src/tensor_ops/attention_reshape/cpu_kernel.rs b/dfdx-core/src/tensor_ops/attention_reshape/cpu_kernel.rs index bebe6883a..cda993c0d 100644 --- a/dfdx-core/src/tensor_ops/attention_reshape/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/attention_reshape/cpu_kernel.rs @@ -13,7 +13,7 @@ impl super::AttentionReshapeKernel for Cpu { Tensor<(Const, Const, usize), E, Self>, Tensor<(Const, usize, Const), E, Self>, ), - Self::Err, + Error, > { let sequence_length = qkv.shape().0; let past_sequence_length = past_key.shape().2; diff --git a/dfdx-core/src/tensor_ops/attention_reshape/cuda_kernel.rs b/dfdx-core/src/tensor_ops/attention_reshape/cuda_kernel.rs index 9f34490ed..1f121e637 100644 --- a/dfdx-core/src/tensor_ops/attention_reshape/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/attention_reshape/cuda_kernel.rs @@ -54,7 +54,7 @@ where Tensor<(Const, Const, usize), E, Self>, Tensor<(Const, usize, Const), E, Self>, ), - Self::Err, + Error, > { if !self.dev.has_func(Self::FN, Self::FN) { self.dev.load_ptx(PTX.into(), Self::FN, &[Self::FN])?; diff --git a/dfdx-core/src/tensor_ops/attention_reshape/mod.rs b/dfdx-core/src/tensor_ops/attention_reshape/mod.rs index b5eeb1b82..b61cc3e4e 100644 --- a/dfdx-core/src/tensor_ops/attention_reshape/mod.rs +++ b/dfdx-core/src/tensor_ops/attention_reshape/mod.rs @@ -57,7 +57,7 @@ pub trait TryAttentionReshape: Storage { qkv: &Tensor<(usize, Const), E, Self>, past_key: &Tensor<(Const, Const, usize), E, Self>, past_value: &Tensor<(Const, usize, Const), E, Self>, - ) -> Result, Self::Err>; + ) -> Result, Error>; } pub trait AttentionReshapeKernel: Storage { @@ -66,7 +66,7 @@ pub trait AttentionReshapeKernel: Storage { qkv: &Tensor<(usize, Const), E, Self>, past_key: &Tensor<(Const, Const, usize), E, Self>, past_value: &Tensor<(Const, usize, Const), E, Self>, - ) -> Result, Self::Err>; + ) -> Result, Error>; } impl> TryAttentionReshape for D { @@ -80,7 +80,7 @@ impl> TryAttentionReshape for D { qkv: &Tensor<(usize, Const), E, Self>, past_key: &Tensor<(Const, Const, usize), E, Self>, past_value: &Tensor<(Const, usize, Const), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let device = qkv.device.clone(); device.forward(qkv, past_key, past_value) } diff --git a/dfdx-core/src/tensor_ops/axpy/cpu_kernel.rs b/dfdx-core/src/tensor_ops/axpy/cpu_kernel.rs index 760d17aa7..c94083528 100644 --- a/dfdx-core/src/tensor_ops/axpy/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/axpy/cpu_kernel.rs @@ -1,13 +1,10 @@ -use crate::{shapes::Dtype, tensor::Cpu}; +use crate::{ + shapes::Dtype, + tensor::{Cpu, Error}, +}; impl super::AxpyKernel for Cpu { - fn forward( - &self, - a: &mut Self::Vec, - alpha: E, - b: &Self::Vec, - beta: E, - ) -> Result<(), Self::Err> { + fn forward(&self, a: &mut Self::Vec, alpha: E, b: &Self::Vec, beta: E) -> Result<(), Error> { for (a_i, b_i) in a.iter_mut().zip(b.iter()) { *a_i = *a_i * alpha + *b_i * beta; } diff --git a/dfdx-core/src/tensor_ops/axpy/cuda_kernel.rs b/dfdx-core/src/tensor_ops/axpy/cuda_kernel.rs index 5a2abf0c9..9d5b7db4e 100644 --- a/dfdx-core/src/tensor_ops/axpy/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/axpy/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ dtypes::*, - tensor::{launch_cfg, Cuda}, + tensor::{launch_cfg, Cuda, Error}, }; use cudarc::driver::{DeviceSlice, LaunchAsync}; @@ -29,13 +29,7 @@ impl super::AxpyKernel for Cuda where Self: HasCudaKernel, { - fn forward( - &self, - a: &mut Self::Vec, - alpha: E, - b: &Self::Vec, - beta: E, - ) -> Result<(), Self::Err> { + fn forward(&self, a: &mut Self::Vec, alpha: E, b: &Self::Vec, beta: E) -> Result<(), Error> { if !self.dev.has_func(Self::FN, Self::FN) { self.dev.load_ptx(PTX_SRC.into(), Self::FN, &[Self::FN])?; } diff --git a/dfdx-core/src/tensor_ops/axpy/mod.rs b/dfdx-core/src/tensor_ops/axpy/mod.rs index ca108c57b..45e6f4653 100644 --- a/dfdx-core/src/tensor_ops/axpy/mod.rs +++ b/dfdx-core/src/tensor_ops/axpy/mod.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Dtype, Shape}, - tensor::{Storage, Tensor}, + tensor::{Error, Storage, Tensor}, }; mod cpu_kernel; @@ -36,7 +36,7 @@ impl> Tensor { alpha: impl Into, b: &Tensor, beta: impl Into, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { assert_eq!(self.shape, b.shape); assert_eq!(self.strides, b.strides, "Strides must be equal for axpy"); self.device.clone().forward( @@ -49,8 +49,7 @@ impl> Tensor { } pub trait AxpyKernel: Storage { - fn forward(&self, a: &mut Self::Vec, alpha: E, b: &Self::Vec, beta: E) - -> Result<(), Self::Err>; + fn forward(&self, a: &mut Self::Vec, alpha: E, b: &Self::Vec, beta: E) -> Result<(), Error>; } #[cfg(test)] diff --git a/dfdx-core/src/tensor_ops/bce/mod.rs b/dfdx-core/src/tensor_ops/bce/mod.rs index 3053f045f..0df0bfeb6 100644 --- a/dfdx-core/src/tensor_ops/bce/mod.rs +++ b/dfdx-core/src/tensor_ops/bce/mod.rs @@ -45,7 +45,10 @@ impl, LTape: Tape> self.try_bce_with_logits(prob).unwrap() } /// See [bce_with_logits] - pub fn try_bce_with_logits(self, prob: Tensor) -> Result + pub fn try_bce_with_logits( + self, + prob: Tensor, + ) -> Result where RTape: Tape, LTape: Merge, diff --git a/dfdx-core/src/tensor_ops/boolean/cpu_kernels.rs b/dfdx-core/src/tensor_ops/boolean/cpu_kernels.rs index ebdab03d7..58a94bb1e 100644 --- a/dfdx-core/src/tensor_ops/boolean/cpu_kernels.rs +++ b/dfdx-core/src/tensor_ops/boolean/cpu_kernels.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Shape, Unit}, - tensor::{cpu::LendingIterator, Cpu, HasErr, Tensor, ZerosTensor}, + tensor::{cpu::LendingIterator, Cpu, Error, Tensor, ZerosTensor}, }; use super::BooleanKernel; @@ -11,7 +11,7 @@ impl Cpu { op: O, lhs: &Tensor, rhs: &Tensor, - ) -> Result, ::Err> { + ) -> Result, Error> { let mut out = self.try_zeros_like(&lhs.shape)?; let mut lhs_iter = lhs.iter(); let mut rhs_iter = rhs.iter(); @@ -24,10 +24,7 @@ impl Cpu { } impl BooleanKernel for Cpu { - fn not( - &self, - inp: &Tensor, - ) -> Result, Self::Err> { + fn not(&self, inp: &Tensor) -> Result, Error> { let mut out = inp.clone(); for x in out.buf_iter_mut() { *x = !*x; @@ -39,7 +36,7 @@ impl BooleanKernel for Cpu { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.eval_binary(|l, r| l && r, lhs, rhs) } @@ -47,7 +44,7 @@ impl BooleanKernel for Cpu { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.eval_binary(|l, r| l || r, lhs, rhs) } @@ -55,7 +52,7 @@ impl BooleanKernel for Cpu { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.eval_binary(|l, r| l ^ r, lhs, rhs) } } diff --git a/dfdx-core/src/tensor_ops/boolean/cuda_kernels.rs b/dfdx-core/src/tensor_ops/boolean/cuda_kernels.rs index e3ecf41a9..f0031ad4c 100644 --- a/dfdx-core/src/tensor_ops/boolean/cuda_kernels.rs +++ b/dfdx-core/src/tensor_ops/boolean/cuda_kernels.rs @@ -1,7 +1,7 @@ use super::BooleanKernel; use crate::{ shapes::Shape, - tensor::{launch_cfg, Cuda, CudaError, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::driver::*; @@ -15,7 +15,7 @@ impl Cuda { fn_name: &str, lhs: &Tensor, rhs: &Tensor, - ) -> Result, CudaError> { + ) -> Result, Error> { if !self.dev.has_func(MODULE_NAME, fn_name) { self.dev .load_ptx(PTX_SRC.into(), MODULE_NAME, &ALL_FN_NAMES)?; @@ -49,10 +49,7 @@ impl Cuda { } impl BooleanKernel for Cuda { - fn not( - &self, - inp: &Tensor, - ) -> Result, Self::Err> { + fn not(&self, inp: &Tensor) -> Result, Error> { if !self.dev.has_func(MODULE_NAME, "boolean_not") { self.dev .load_ptx(PTX_SRC.into(), MODULE_NAME, &ALL_FN_NAMES)?; @@ -77,7 +74,7 @@ impl BooleanKernel for Cuda { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.call_binary("boolean_and", lhs, rhs) } @@ -85,7 +82,7 @@ impl BooleanKernel for Cuda { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.call_binary("boolean_or", lhs, rhs) } @@ -93,7 +90,7 @@ impl BooleanKernel for Cuda { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.call_binary("boolean_xor", lhs, rhs) } } diff --git a/dfdx-core/src/tensor_ops/boolean/mod.rs b/dfdx-core/src/tensor_ops/boolean/mod.rs index 51c9c5def..bb8cedf42 100644 --- a/dfdx-core/src/tensor_ops/boolean/mod.rs +++ b/dfdx-core/src/tensor_ops/boolean/mod.rs @@ -6,7 +6,7 @@ mod cuda_kernels; use crate::{ prelude::{OnesTensor, Tensor, ZerosTensor}, shapes::*, - tensor::Storage, + tensor::{Error, Storage}, }; use std::ops::{BitAnd, BitOr, BitXor, Not}; @@ -14,34 +14,31 @@ use std::ops::{BitAnd, BitOr, BitXor, Not}; use super::Device; pub trait BooleanKernel: Storage + OnesTensor + ZerosTensor { - fn not( - &self, - inp: &Tensor, - ) -> Result, Self::Err>; + fn not(&self, inp: &Tensor) -> Result, Error>; fn and( &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn or( &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn xor( &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; } fn scalar_and( lhs: &Tensor, rhs: bool, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { if rhs { Ok(lhs.clone()) } else { @@ -52,7 +49,7 @@ fn scalar_and( fn scalar_or( lhs: &Tensor, rhs: bool, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { if rhs { lhs.device.try_ones_like(lhs) } else { @@ -63,7 +60,7 @@ fn scalar_or( fn scalar_xor( lhs: &Tensor, rhs: bool, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { if rhs { Ok(lhs.device.not(lhs)?) } else { diff --git a/dfdx-core/src/tensor_ops/broadcast_to.rs b/dfdx-core/src/tensor_ops/broadcast_to.rs index 90dfa93a9..789639c73 100644 --- a/dfdx-core/src/tensor_ops/broadcast_to.rs +++ b/dfdx-core/src/tensor_ops/broadcast_to.rs @@ -23,7 +23,7 @@ use crate::{shapes::*, tensor::*}; /// // It's ambiguous what axes to broadcast here - explicitly say axes 0 and 2 /// let _: Tensor, _, _> = a.clone().broadcast::<_, Axes2<0, 2>>(); /// ``` -pub trait BroadcastTo: HasErr + HasShape { +pub trait BroadcastTo: Sized + HasShape { /// Broadcast into shape `Dst` along axes `Ax`. fn broadcast(self) -> Self::WithShape where @@ -33,7 +33,7 @@ pub trait BroadcastTo: HasErr + HasShape { .unwrap() } /// Fallible version of [BroadcastTo::broadcast] - fn try_broadcast(self) -> Result, Self::Err> + fn try_broadcast(self) -> Result, Error> where Self::Shape: BroadcastShapeTo, { @@ -50,7 +50,7 @@ pub trait BroadcastTo: HasErr + HasShape { fn try_broadcast_like( self, dst: &Dst, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: BroadcastShapeTo; } @@ -59,7 +59,7 @@ impl, T: Tape> BroadcastTo for Tensor( self, dst: &Dst, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: BroadcastShapeTo, { diff --git a/dfdx-core/src/tensor_ops/choose/cpu_kernel.rs b/dfdx-core/src/tensor_ops/choose/cpu_kernel.rs index 2d7b15f27..068579311 100644 --- a/dfdx-core/src/tensor_ops/choose/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/choose/cpu_kernel.rs @@ -2,7 +2,7 @@ use crate::{ shapes::{Dtype, Shape}, tensor::{ cpu::{LendingIterator, NdIndex}, - Cpu, Storage, Tensor, ZerosTensor, + Cpu, Error, Storage, Tensor, ZerosTensor, }, }; @@ -12,7 +12,7 @@ impl super::ChooseKernel for Cpu { cond: &Tensor, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut out = self.try_zeros_like(&lhs.shape)?; let mut cond_iter = cond.iter(); let mut lhs_iter = lhs.iter(); @@ -36,7 +36,7 @@ impl super::ChooseKernel for Cpu { rhs: &Tensor, grad_rhs: &mut >::Vec, grad_out: &>::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut lhs_idx = NdIndex::new(lhs.shape, lhs.strides); let mut rhs_idx = NdIndex::new(rhs.shape, rhs.strides); let mut out_idx = NdIndex::new(lhs.shape, lhs.shape.strides()); diff --git a/dfdx-core/src/tensor_ops/choose/cuda_kernel.rs b/dfdx-core/src/tensor_ops/choose/cuda_kernel.rs index ed29149bb..e99898c61 100644 --- a/dfdx-core/src/tensor_ops/choose/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/choose/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Storage, Tensor}, + tensor::{launch_cfg, Cuda, Error, Storage, Tensor}, }; use cudarc::driver::{CudaSlice, LaunchAsync}; @@ -43,7 +43,7 @@ where cond: &Tensor, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(Self::MOD, Self::FNS[0]) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?; } @@ -85,7 +85,7 @@ where rhs: &Tensor, grad_rhs: &mut >::Vec, grad_out: &>::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let bwd_fn = self.dev.get_func(Self::MOD, Self::FNS[1]).unwrap(); let numel = cond.shape.num_elements(); diff --git a/dfdx-core/src/tensor_ops/choose/mod.rs b/dfdx-core/src/tensor_ops/choose/mod.rs index a1861e448..a82a5c79c 100644 --- a/dfdx-core/src/tensor_ops/choose/mod.rs +++ b/dfdx-core/src/tensor_ops/choose/mod.rs @@ -5,7 +5,7 @@ mod cuda_kernel; use crate::{ shapes::{Dtype, HasShape, Shape}, - tensor::{HasErr, Merge, PutTape, SplitTape, Storage, Tape, Tensor}, + tensor::{Error, Merge, PutTape, SplitTape, Storage, Tape, Tensor}, }; pub trait ChooseKernel: Storage + Storage { @@ -14,7 +14,7 @@ pub trait ChooseKernel: Storage + Storage { cond: &Tensor, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, @@ -24,7 +24,7 @@ pub trait ChooseKernel: Storage + Storage { rhs: &Tensor, grad_rhs: &mut >::Vec, grad_out: &>::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Choose values from two tensors using a boolean mask. Equivalent to `torch.where` from pytorch. @@ -38,7 +38,7 @@ pub trait ChooseKernel: Storage + Storage { /// let c = cond.choose(a, b); /// assert_eq!(c.array(), [1.0, -2.0, 3.0]); /// ``` -pub trait ChooseFrom: HasErr { +pub trait ChooseFrom: Sized { type Output; /// Construct a new tensor, where the output tensor contains the elements of lhs where self is @@ -48,7 +48,7 @@ pub trait ChooseFrom: HasErr { } /// Fallible version of choose - fn try_choose(self, lhs: Lhs, rhs: Rhs) -> Result; + fn try_choose(self, lhs: Lhs, rhs: Rhs) -> Result; } impl< @@ -65,7 +65,7 @@ impl< self, lhs: Tensor, rhs: Tensor, - ) -> Result { + ) -> Result { assert_eq!(self.shape(), lhs.shape()); assert_eq!(lhs.shape(), rhs.shape()); diff --git a/dfdx-core/src/tensor_ops/clamp/mod.rs b/dfdx-core/src/tensor_ops/clamp/mod.rs index b23f6a035..1054d0ffa 100644 --- a/dfdx-core/src/tensor_ops/clamp/mod.rs +++ b/dfdx-core/src/tensor_ops/clamp/mod.rs @@ -37,7 +37,11 @@ impl, E>, T: Tape> Ten self.try_clamp(min, max).unwrap() } /// See [clamp] - pub fn try_clamp(self, min: impl Into, max: impl Into) -> Result { + pub fn try_clamp( + self, + min: impl Into, + max: impl Into, + ) -> Result { try_unary_op( ClampKernelOp { min: E::from_f64(min.into()).unwrap(), diff --git a/dfdx-core/src/tensor_ops/cmp/cpu_kernels.rs b/dfdx-core/src/tensor_ops/cmp/cpu_kernels.rs index c42ddb543..8a22615cc 100644 --- a/dfdx-core/src/tensor_ops/cmp/cpu_kernels.rs +++ b/dfdx-core/src/tensor_ops/cmp/cpu_kernels.rs @@ -2,7 +2,7 @@ use crate::{ shapes::{Shape, Unit}, tensor::{ cpu::{Cpu, LendingIterator}, - Tensor, ZerosTensor, + Error, Tensor, ZerosTensor, }, }; @@ -20,7 +20,7 @@ impl, E: Unit> CmpKernel for Cpu { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut out: Tensor = self.try_zeros_like(&lhs.shape)?; let mut lhs_iter = lhs.iter(); let mut rhs_iter = rhs.iter(); @@ -37,7 +37,7 @@ impl, E: Unit> ScalarCmpKernel for Cpu { &self, lhs: &Tensor, scalar: E, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut out: Tensor = self.try_zeros_like(&lhs.shape)?; let mut lhs_iter = lhs.iter(); let mut out_iter = out.iter_mut(); diff --git a/dfdx-core/src/tensor_ops/cmp/cuda_kernels.rs b/dfdx-core/src/tensor_ops/cmp/cuda_kernels.rs index 31d0d1196..b4a43cb8e 100644 --- a/dfdx-core/src/tensor_ops/cmp/cuda_kernels.rs +++ b/dfdx-core/src/tensor_ops/cmp/cuda_kernels.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::Shape, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::driver::{CudaSlice, LaunchAsync}; @@ -39,7 +39,7 @@ impl> CmpKernel for Cuda { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(Op::MODULE_NAME, Op::FWD_FN_NAME) { self.dev .load_ptx(Op::PTX_SRC.into(), Op::MODULE_NAME, &[Op::FWD_FN_NAME])?; @@ -80,7 +80,7 @@ impl> ScalarCmpKernel for Cuda { &self, lhs: &Tensor, scalar: E, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(Op::MODULE_NAME, Op::FWD_FN_NAME) { self.dev .load_ptx(Op::PTX_SRC.into(), Op::MODULE_NAME, &[Op::FWD_FN_NAME])?; diff --git a/dfdx-core/src/tensor_ops/cmp/mod.rs b/dfdx-core/src/tensor_ops/cmp/mod.rs index cd1c3cc08..cf5feed3f 100644 --- a/dfdx-core/src/tensor_ops/cmp/mod.rs +++ b/dfdx-core/src/tensor_ops/cmp/mod.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{HasShape, Shape}, - tensor::{HasErr, NoneTape, Storage, Tape, Tensor}, + tensor::{Error, NoneTape, Storage, Tape, Tensor}, }; mod cpu_kernels; @@ -12,13 +12,13 @@ pub trait CmpKernel: Storage + Storage { &self, lhs: &Tensor, rhs: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; } fn try_cmp_op, T: Tape>( lhs: &Tensor, rhs: &Tensor, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { assert_eq!(lhs.shape(), rhs.shape()); lhs.device.forward(lhs, rhs) } @@ -28,13 +28,13 @@ pub trait ScalarCmpKernel: Storage + Storage { &self, tensor: &Tensor, scalar: E, - ) -> Result, Self::Err>; + ) -> Result, Error>; } fn try_scalar_cmp_op, T: Tape>( tensor: &Tensor, scalar: E, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { tensor.device.forward(tensor, scalar) } @@ -198,14 +198,14 @@ pub fn le, T: Tape>( // Macro to reduce boilerplate of implementing comparison methods on Tensor. macro_rules! impl_cmp_kernel_op { ($TraitName:tt, $FnName:tt, $TryFnName:tt, $KernelOp:tt, $doc:expr, $ScalarFnName:tt, $TryScalarFnName:tt) => { - pub trait $TraitName: HasErr { + pub trait $TraitName { type Output; #[doc = $doc] fn $FnName(&self, rhs: Rhs) -> Self::Output { self.$TryFnName(rhs).unwrap() } #[doc = $doc] - fn $TryFnName(&self, rhs: Rhs) -> Result; + fn $TryFnName(&self, rhs: Rhs) -> Result; } impl, T: Tape> $TraitName<&Self> @@ -213,7 +213,7 @@ macro_rules! impl_cmp_kernel_op { { type Output = Tensor; #[doc = $doc] - fn $TryFnName(&self, other: &Self) -> Result { + fn $TryFnName(&self, other: &Self) -> Result { try_cmp_op(self, other) } } @@ -223,7 +223,7 @@ macro_rules! impl_cmp_kernel_op { { type Output = Tensor; #[doc = $doc] - fn $TryFnName(&self, other: E) -> Result { + fn $TryFnName(&self, other: E) -> Result { try_scalar_cmp_op(self, other) } } @@ -234,7 +234,7 @@ macro_rules! impl_cmp_kernel_op { { type Output = Tensor; #[doc = $doc] - fn $TryFnName(&self, other: f32) -> Result { + fn $TryFnName(&self, other: f32) -> Result { try_scalar_cmp_op(self, half::f16::from_f32(other)) } } @@ -248,7 +248,7 @@ macro_rules! impl_cmp_kernel_op { { type Output = Tensor; #[doc = $doc] - fn $TryFnName(&self, other: f32) -> Result { + fn $TryFnName(&self, other: f32) -> Result { try_scalar_cmp_op(self, crate::dtypes::AMP(half::f16::from_f32(other))) } } @@ -265,7 +265,7 @@ macro_rules! impl_cmp_kernel_op { pub fn $TryScalarFnName( &self, other: E, - ) -> Result, D::Err> { + ) -> Result, crate::tensor::Error> { try_scalar_cmp_op(self, other) } } diff --git a/dfdx-core/src/tensor_ops/concat/cpu_kernel.rs b/dfdx-core/src/tensor_ops/concat/cpu_kernel.rs index 19b1dd5f7..63d864bde 100644 --- a/dfdx-core/src/tensor_ops/concat/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/concat/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Dtype, Shape}, - tensor::{unique_id, Cpu, Tensor}, + tensor::{unique_id, Cpu, Error, Tensor}, }; impl super::ConcatKernel for Cpu { @@ -8,7 +8,7 @@ impl super::ConcatKernel for Cpu { &self, a: &Tensor, b: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where A: super::ConcatShape, { @@ -45,7 +45,7 @@ impl super::ConcatKernel for Cpu { grad_a: &mut Self::Vec, grad_b: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut offset = 0; for ga in grad_a.iter_mut() { *ga += grad_out[offset]; diff --git a/dfdx-core/src/tensor_ops/concat/cuda_kernel.rs b/dfdx-core/src/tensor_ops/concat/cuda_kernel.rs index 147c6273b..8da4df742 100644 --- a/dfdx-core/src/tensor_ops/concat/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/concat/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::{ driver::{DeviceSlice, LaunchAsync}, @@ -13,7 +13,7 @@ impl super::ConcatKernel for Cuda { &self, a: &Tensor, b: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where A: super::ConcatShape, { @@ -33,7 +33,7 @@ impl super::ConcatKernel for Cuda { grad_a: &mut Self::Vec, grad_b: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let module_name = std::format!("concat_bwd_{}", E::NAME); if !self.dev.has_func(&module_name, "concat_bwd") { let src = BWD_KERNEL.replace("$Ty", E::NAME); diff --git a/dfdx-core/src/tensor_ops/concat/mod.rs b/dfdx-core/src/tensor_ops/concat/mod.rs index 0ad3a9669..da6cdfee7 100644 --- a/dfdx-core/src/tensor_ops/concat/mod.rs +++ b/dfdx-core/src/tensor_ops/concat/mod.rs @@ -27,7 +27,7 @@ mod cuda_kernel; /// assert_eq!(c.shape().0, 6); /// ``` #[deprecated = "Use TryConcatAlong instead"] -pub trait TryConcat: HasErr { +pub trait TryConcat: Sized { type Output; /// Concatenate two tensors along the first dimension. @@ -41,7 +41,7 @@ pub trait TryConcat: HasErr { /// Fallible version of [TryConcat::concat]. #[deprecated = "Use TryConcatAlong::try_concat_along instead"] #[allow(deprecated)] - fn try_concat(self, rhs: Rhs) -> Result; + fn try_concat(self, rhs: Rhs) -> Result; } #[allow(deprecated)] @@ -54,7 +54,7 @@ where { type Output = Tensor; #[allow(deprecated)] - fn try_concat(self, rhs: Tensor) -> Result { + fn try_concat(self, rhs: Tensor) -> Result { assert_eq!( self.strides, self.shape.strides(), @@ -89,7 +89,7 @@ pub trait ConcatKernel: Storage { &self, a: &Tensor, b: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where A: ConcatShape; fn backward( @@ -97,7 +97,7 @@ pub trait ConcatKernel: Storage { grad_a: &mut Self::Vec, grad_b: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } pub trait ConcatShape: Shape { diff --git a/dfdx-core/src/tensor_ops/concat_along/cpu_kernel.rs b/dfdx-core/src/tensor_ops/concat_along/cpu_kernel.rs index 7404b9aff..e6ab2eb20 100644 --- a/dfdx-core/src/tensor_ops/concat_along/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/concat_along/cpu_kernel.rs @@ -10,7 +10,7 @@ impl super::ConcatAlongKernel for Cpu { a: &Tensor, b: &Tensor, c: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut a_idx = NdIndex::new(a.shape, a.strides); let mut b_idx = NdIndex::new(b.shape, b.strides); @@ -44,7 +44,7 @@ impl super::ConcatAlongKernel for Cpu { b: &GhostTensor, grad_b: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut a_idx = NdIndex::new(a.shape, a.strides); let mut b_idx = NdIndex::new(b.shape, b.strides); diff --git a/dfdx-core/src/tensor_ops/concat_along/cuda_kernel.rs b/dfdx-core/src/tensor_ops/concat_along/cuda_kernel.rs index 7c0d1247c..c779cadf2 100644 --- a/dfdx-core/src/tensor_ops/concat_along/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/concat_along/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::*, - tensor::{launch_cfg, Cuda, GhostTensor, Tensor}, + tensor::{launch_cfg, Cuda, Error, GhostTensor, Tensor}, }; use cudarc::{ driver::{DeviceSlice, LaunchAsync}, @@ -15,7 +15,7 @@ impl super::ConcatAlongKernel for Cuda { a: &Tensor, b: &Tensor, c: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let module_name = std::format!("concat_{}", E::NAME); if !self.dev.has_func(&module_name, "fwd") { let src = KERNEL.replace("$Ty", E::NAME); @@ -67,7 +67,7 @@ impl super::ConcatAlongKernel for Cuda { b: &GhostTensor, grad_b: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let module_name = std::format!("concat_{}", E::NAME); let bwd = self.dev.get_func(&module_name, "bwd").unwrap(); let cfg = launch_cfg::<128>(grad_out.data.len() as u32); diff --git a/dfdx-core/src/tensor_ops/concat_along/mod.rs b/dfdx-core/src/tensor_ops/concat_along/mod.rs index 28282419f..0c796d6e6 100644 --- a/dfdx-core/src/tensor_ops/concat_along/mod.rs +++ b/dfdx-core/src/tensor_ops/concat_along/mod.rs @@ -48,13 +48,13 @@ mod cuda_kernel; /// ``` pub trait TryConcatAlong: Sized { type Output; - type Error: std::fmt::Debug; + /// Concatenates self along the given axis. fn concat_along(self, ax: Ax) -> Self::Output { self.try_concat_along(ax).unwrap() } /// Fallibly concatenates self along the given axis. - fn try_concat_along(self, ax: Ax) -> Result; + fn try_concat_along(self, ax: Ax) -> Result; } pub trait ConcatAlongKernel: Storage { @@ -64,7 +64,7 @@ pub trait ConcatAlongKernel: Storage { a: &Tensor, b: &Tensor, c: &mut Tensor, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; fn backward( &self, @@ -74,7 +74,7 @@ pub trait ConcatAlongKernel: Storage { b: &GhostTensor, grad_b: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl, R: Tape> TryConcatAlong @@ -89,8 +89,8 @@ where T: Merge, { type Output = Tensor<<(A, B) as TryConcatAlong>::Output, E, D, T>; - type Error = D::Err; - fn try_concat_along(self, ax: Ax) -> Result { + + fn try_concat_along(self, ax: Ax) -> Result { let (lhs, rhs) = self; let out_shape = (*lhs.shape(), *rhs.shape()).concat_along(ax); @@ -135,8 +135,8 @@ macro_rules! impl_concat { >::Output, $($Tail, )* ); - type Error = std::convert::Infallible; - fn try_concat_along(self, _: Axis<$Ax>) -> Result { + + fn try_concat_along(self, _: Axis<$Ax>) -> Result { let (lhs, rhs) = self; let lhs_dims = lhs.concrete(); let rhs_dims = rhs.concrete(); diff --git a/dfdx-core/src/tensor_ops/conv1d/cpu_kernel.rs b/dfdx-core/src/tensor_ops/conv1d/cpu_kernel.rs index 4383c9eb2..97798dd60 100644 --- a/dfdx-core/src/tensor_ops/conv1d/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/conv1d/cpu_kernel.rs @@ -34,7 +34,7 @@ impl Cpu { filters: &[E], out: &mut [E], buf: &mut [E], - ) -> Result<(), CpuError> + ) -> Result<(), Error> where Self: MatMulImpl, { @@ -85,7 +85,7 @@ impl Cpu { grad_filters_tr: &mut [E], grad_out: &[E], buf: &mut [E], - ) -> Result<(), CpuError> + ) -> Result<(), Error> where Self: MatMulImpl, { @@ -150,7 +150,7 @@ impl Conv1DKernel for Cpu where Self: MatMulImpl, { - fn alloc(&self, s: S) -> Result, Self::Err> { + fn alloc(&self, s: S) -> Result, Error> { self.try_zeros_like(&s) } @@ -160,7 +160,7 @@ where lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches = (op.chan_in, op.kernel, op.l_out); let mut patches = self.try_alloc_zeros::(patches.num_elements())?; let [lstride, ostride] = match L::NUM_DIMS { @@ -192,7 +192,7 @@ where grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let f_tr_shape = [ op.groups, op.chan_in / op.groups, diff --git a/dfdx-core/src/tensor_ops/conv1d/cuda_kernel.rs b/dfdx-core/src/tensor_ops/conv1d/cuda_kernel.rs index 4443d9da0..45779e26d 100644 --- a/dfdx-core/src/tensor_ops/conv1d/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/conv1d/cuda_kernel.rs @@ -4,7 +4,7 @@ use cudarc::driver::{DeviceRepr, LaunchAsync, ValidAsZeroBits}; use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor, Tensorlike}, + tensor::{launch_cfg, Cuda, Error, Tensor, Tensorlike}, }; use std::sync::Arc; @@ -73,7 +73,7 @@ where Self: HasCudaKernel, CudaBlas: Gemm, { - fn alloc(&self, shape: S) -> Result, Self::Err> { + fn alloc(&self, shape: S) -> Result, Error> { let data = unsafe { self.alloc_empty::(shape.num_elements()) }?; Ok(self.build_tensor(shape, shape.strides(), data)) } @@ -83,7 +83,7 @@ where img: &Tensor, fil: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FNS[0]) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?; } @@ -151,7 +151,7 @@ where grad_rhs: &mut Self::Vec, _: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches_item_numel = op.chan_out * op.kernel * op.l_in; let patches_numel = op.batch * patches_item_numel; let filters_numel = diff --git a/dfdx-core/src/tensor_ops/conv1d/mod.rs b/dfdx-core/src/tensor_ops/conv1d/mod.rs index 25bdddbee..db51c02bf 100644 --- a/dfdx-core/src/tensor_ops/conv1d/mod.rs +++ b/dfdx-core/src/tensor_ops/conv1d/mod.rs @@ -23,7 +23,7 @@ pub(super) struct Conv1DOp { } pub(super) trait Conv1DKernel: Storage { - fn alloc(&self, s: S) -> Result, Self::Err>; + fn alloc(&self, s: S) -> Result, Error>; fn forward( &self, @@ -31,7 +31,7 @@ pub(super) trait Conv1DKernel: Storage { lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; #[allow(clippy::too_many_arguments)] fn backward( @@ -43,7 +43,7 @@ pub(super) trait Conv1DKernel: Storage { grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Applies a 1d convolution to a tensor. @@ -86,7 +86,6 @@ pub(super) trait Conv1DKernel: Storage { /// ``` pub trait TryConv1D: Sized { type Convolved; - type Error: std::fmt::Debug; /// Applies a 1D convolution to the input tensor. fn conv1d( @@ -106,7 +105,7 @@ pub trait TryConv1D: Sized { padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result; + ) -> Result; } impl< @@ -122,14 +121,13 @@ where Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>: Sized, { type Convolved = Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>; - type Error = std::convert::Infallible; fn try_conv1d( self, _: Const, _: Const, _: Const, _: Groups, - ) -> Result { + ) -> Result { Ok(Const) } } @@ -138,14 +136,13 @@ impl TryConv1D for (usize, Kernel) { type Convolved = usize; - type Error = std::convert::Infallible; fn try_conv1d( self, stride: Stride, padding: Padding, dilation: Dilation, _: Groups, - ) -> Result { + ) -> Result { let (dim, kernel) = self; Ok((dim + 2 * padding.size() - 1) .checked_sub(dilation.size() * (kernel.size() - 1)) @@ -187,15 +184,13 @@ where D, T, >; - type Error = D::Err; - fn try_conv1d( self, stride: Stride, padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result { + ) -> Result { let (img, filters) = self; let (inp_chan, l) = img.shape; let img = img.try_reshape_like(&(Const::<1>, inp_chan, l))?; @@ -239,15 +234,13 @@ where D, T, >; - type Error = D::Err; - fn try_conv1d( self, stride: Stride, padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result { + ) -> Result { let (img, filters) = self; assert_eq!(img.shape.1.size(), filters.shape.1.size() * groups.size()); let (batch, inp_chan, l) = img.shape; diff --git a/dfdx-core/src/tensor_ops/conv2d/cpu_kernel.rs b/dfdx-core/src/tensor_ops/conv2d/cpu_kernel.rs index 5edf859b1..f058fc422 100644 --- a/dfdx-core/src/tensor_ops/conv2d/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/conv2d/cpu_kernel.rs @@ -48,7 +48,7 @@ impl Cpu { filters: &[E], out: &mut [E], buf: &mut [E], - ) -> Result<(), CpuError> + ) -> Result<(), Error> where Self: MatMulImpl, { @@ -105,7 +105,7 @@ impl Cpu { grad_filters_tr: &mut [E], grad_out: &[E], buf: &mut [E], - ) -> Result<(), CpuError> + ) -> Result<(), Error> where Self: MatMulImpl, { @@ -175,7 +175,7 @@ impl Conv2DKernel for Cpu where Self: MatMulImpl, { - fn alloc(&self, s: S) -> Result, Self::Err> { + fn alloc(&self, s: S) -> Result, Error> { self.try_zeros_like(&s) } @@ -185,7 +185,7 @@ where lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches = (op.chan_in, op.kernel, op.kernel, op.h_out, op.w_out); let mut patches = self.try_alloc_zeros::(patches.num_elements())?; let [lstride, ostride] = match L::NUM_DIMS { @@ -217,7 +217,7 @@ where grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let f_tr_shape = [ op.groups, op.chan_in / op.groups, diff --git a/dfdx-core/src/tensor_ops/conv2d/cuda_kernel.rs b/dfdx-core/src/tensor_ops/conv2d/cuda_kernel.rs index 4836ed65c..cb38ad192 100644 --- a/dfdx-core/src/tensor_ops/conv2d/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/conv2d/cuda_kernel.rs @@ -4,7 +4,7 @@ use cudarc::driver::{DeviceRepr, LaunchAsync, ValidAsZeroBits}; use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor, Tensorlike}, + tensor::{launch_cfg, Cuda, Error, Tensor, Tensorlike}, }; use std::sync::Arc; @@ -73,7 +73,7 @@ where Self: HasCudaKernel, CudaBlas: Gemm, { - fn alloc(&self, shape: S) -> Result, Self::Err> { + fn alloc(&self, shape: S) -> Result, Error> { let data = unsafe { self.alloc_empty::(shape.num_elements()) }?; Ok(self.build_tensor(shape, shape.strides(), data)) } @@ -83,7 +83,7 @@ where img: &Tensor, fil: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FNS[0]) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?; } @@ -152,7 +152,7 @@ where grad_rhs: &mut Self::Vec, _: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches_item_numel = op.chan_out * op.kernel * op.kernel * op.h_in * op.w_in; let patches_numel = op.batch * patches_item_numel; let filters_numel = op.groups diff --git a/dfdx-core/src/tensor_ops/conv2d/cudnn_kernel.rs b/dfdx-core/src/tensor_ops/conv2d/cudnn_kernel.rs index 8e614e5eb..872740243 100644 --- a/dfdx-core/src/tensor_ops/conv2d/cudnn_kernel.rs +++ b/dfdx-core/src/tensor_ops/conv2d/cudnn_kernel.rs @@ -4,7 +4,7 @@ use cudarc::driver::DeviceSlice; use crate::{ dtypes::*, shapes::*, - tensor::{Cuda, Tensor, Tensorlike}, + tensor::{Cuda, Error, Tensor, Tensorlike}, }; use std::sync::Arc; @@ -29,7 +29,7 @@ impl super::Conv2DKernel for Cuda where Self: HasCudnnKernel, { - fn alloc(&self, shape: S) -> Result, Self::Err> { + fn alloc(&self, shape: S) -> Result, Error> { let data = unsafe { self.alloc_empty::(shape.num_elements()) }?; Ok(self.build_tensor(shape, shape.strides(), data)) } @@ -39,7 +39,7 @@ where lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut conv = self.cudnn.create_conv2d::( [op.padding as i32, op.padding as i32], [op.stride as i32, op.stride as i32], @@ -97,7 +97,7 @@ where grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut conv = self.cudnn.create_conv2d::( [op.padding as i32, op.padding as i32], [op.stride as i32, op.stride as i32], diff --git a/dfdx-core/src/tensor_ops/conv2d/mod.rs b/dfdx-core/src/tensor_ops/conv2d/mod.rs index c61b4cb7d..c5be96945 100644 --- a/dfdx-core/src/tensor_ops/conv2d/mod.rs +++ b/dfdx-core/src/tensor_ops/conv2d/mod.rs @@ -29,7 +29,7 @@ pub(super) struct Conv2DOp { } pub(super) trait Conv2DKernel: Storage { - fn alloc(&self, s: S) -> Result, Self::Err>; + fn alloc(&self, s: S) -> Result, Error>; fn forward( &self, @@ -37,7 +37,7 @@ pub(super) trait Conv2DKernel: Storage { lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; #[allow(clippy::too_many_arguments)] fn backward( @@ -49,7 +49,7 @@ pub(super) trait Conv2DKernel: Storage { grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Apply the 2d convolution to a tensor. @@ -94,7 +94,6 @@ pub(super) trait Conv2DKernel: Storage { /// ``` pub trait TryConv2D: Sized { type Convolved; - type Error: std::fmt::Debug; /// Applies a 2D convolution to the input tensor. fn conv2d( @@ -114,7 +113,7 @@ pub trait TryConv2D: Sized { padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result; + ) -> Result; } impl< @@ -130,14 +129,13 @@ where Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>: Sized, { type Convolved = Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>; - type Error = std::convert::Infallible; fn try_conv2d( self, _: Const, _: Const, _: Const, _: Groups, - ) -> Result { + ) -> Result { Ok(Const) } } @@ -146,14 +144,13 @@ impl TryConv2D for (usize, Kernel) { type Convolved = usize; - type Error = std::convert::Infallible; fn try_conv2d( self, stride: Stride, padding: Padding, dilation: Dilation, _: Groups, - ) -> Result { + ) -> Result { let (dim, kernel) = self; Ok((dim + 2 * padding.size() - 1) .checked_sub(dilation.size() * (kernel.size() - 1)) @@ -208,7 +205,6 @@ where D, T, >; - type Error = D::Err; fn try_conv2d( self, @@ -216,7 +212,7 @@ where padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result { + ) -> Result { let (img, filters) = self; let (inp_chan, h, w) = img.shape; let img = img.try_reshape_like(&(Const::<1>, inp_chan, h, w))?; @@ -273,7 +269,6 @@ where D, T, >; - type Error = D::Err; fn try_conv2d( self, @@ -281,7 +276,7 @@ where padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result { + ) -> Result { let (img, filters) = self; assert_eq!(img.shape.1.size(), filters.shape.1.size() * groups.size()); assert_eq!(filters.shape.2, filters.shape.3); diff --git a/dfdx-core/src/tensor_ops/convtrans2d/cpu_kernel.rs b/dfdx-core/src/tensor_ops/convtrans2d/cpu_kernel.rs index f61943f84..d8e842248 100644 --- a/dfdx-core/src/tensor_ops/convtrans2d/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/convtrans2d/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::prelude::Tensorlike; use crate::shapes::{Dtype, Shape}; -use crate::tensor::{cpu::*, Tensor, ZerosTensor}; +use crate::tensor::{cpu::*, Error, Tensor, ZerosTensor}; use crate::tensor_ops::matmul::cpu_kernel::MatMulImpl; use std::sync::Arc; @@ -27,7 +27,7 @@ impl Cpu { filters_tr: &[E], out: &mut [E], buf: &mut [E], - ) -> Result<(), CpuError> + ) -> Result<(), Error> where Self: MatMulImpl, { @@ -107,7 +107,7 @@ impl Cpu { grad_filters: &mut [E], grad_out: &[E], buf: &mut [E], - ) -> Result<(), CpuError> + ) -> Result<(), Error> where Self: MatMulImpl, { @@ -179,7 +179,7 @@ impl ConvTrans2DKernel for Cpu where Self: MatMulImpl, { - fn alloc(&self, s: S) -> Result, Self::Err> { + fn alloc(&self, s: S) -> Result, Error> { self.try_zeros_like(&s) } @@ -189,7 +189,7 @@ where lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches = (op.chan_in, op.kernel, op.kernel, op.h_out, op.w_out); let mut patches = self.try_alloc_zeros::(patches.num_elements())?; let f_tr_shape = [ @@ -242,7 +242,7 @@ where grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches_shape = [op.chan_out, op.kernel, op.kernel, op.h_in, op.w_in]; let mut patches = self.try_alloc_zeros::(patches_shape.num_elements())?; diff --git a/dfdx-core/src/tensor_ops/convtrans2d/cuda_kernel.rs b/dfdx-core/src/tensor_ops/convtrans2d/cuda_kernel.rs index 6ddcd10c7..8ea312f95 100644 --- a/dfdx-core/src/tensor_ops/convtrans2d/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/convtrans2d/cuda_kernel.rs @@ -4,7 +4,7 @@ use cudarc::driver::{DeviceRepr, LaunchAsync, ValidAsZeroBits}; use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor, Tensorlike}, + tensor::{launch_cfg, Cuda, Error, Tensor, Tensorlike}, }; use std::sync::Arc; @@ -69,7 +69,7 @@ where Self: HasCudaKernel, CudaBlas: Gemm, { - fn alloc(&self, shape: S) -> Result, Self::Err> { + fn alloc(&self, shape: S) -> Result, Error> { let data = unsafe { self.alloc_empty::(shape.num_elements()) }?; Ok(self.build_tensor(shape, shape.strides(), data)) } @@ -80,7 +80,7 @@ where lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FNS[0]) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?; } @@ -160,7 +160,7 @@ where grad_rhs: &mut Self::Vec, _: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let patches_numel = op.batch * op.chan_out * op.kernel * op.kernel * op.h_in * op.w_in; let mut patches = unsafe { self.get_workspace::(patches_numel) }?; diff --git a/dfdx-core/src/tensor_ops/convtrans2d/mod.rs b/dfdx-core/src/tensor_ops/convtrans2d/mod.rs index dc56a2fc8..761ab4915 100644 --- a/dfdx-core/src/tensor_ops/convtrans2d/mod.rs +++ b/dfdx-core/src/tensor_ops/convtrans2d/mod.rs @@ -28,7 +28,7 @@ pub(super) struct ConvTrans2DOp { } pub(super) trait ConvTrans2DKernel: Storage { - fn alloc(&self, s: S) -> Result, Self::Err>; + fn alloc(&self, s: S) -> Result, Error>; fn forward( &self, @@ -36,7 +36,7 @@ pub(super) trait ConvTrans2DKernel: Storage { lhs: &Tensor, rhs: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; #[allow(clippy::too_many_arguments)] fn backward( @@ -48,12 +48,11 @@ pub(super) trait ConvTrans2DKernel: Storage { grad_rhs: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } pub trait TryConvTrans2D: Sized { type Convolved; - type Error: std::fmt::Debug; /// Applies a 2D convolution to the input tensor. fn convtrans2d( @@ -74,7 +73,7 @@ pub trait TryConvTrans2D: Sized { padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result; + ) -> Result; } impl< @@ -90,7 +89,6 @@ where Const<{ (DIM - 1) * STRIDE - 2 * PADDING + DILATION * (KERNEL - 1) + 1 }>: Sized, { type Convolved = Const<{ (DIM - 1) * STRIDE - 2 * PADDING + DILATION * (KERNEL - 1) + 1 }>; - type Error = std::convert::Infallible; fn try_convtrans2d( self, @@ -98,7 +96,7 @@ where _: Const, _: Const, _: Groups, - ) -> Result { + ) -> Result { Ok(Const) } } @@ -107,7 +105,6 @@ impl TryConvTrans2D for (usize, Kernel) { type Convolved = usize; - type Error = std::convert::Infallible; fn try_convtrans2d( self, @@ -115,7 +112,7 @@ impl padding: Padding, dilation: Dilation, _: Groups, - ) -> Result { + ) -> Result { let (dim, kernel) = self; Ok( ((dim - 1) * stride.size() + dilation.size() * (kernel.size() - 1) + 1) @@ -161,7 +158,6 @@ where D, T, >; - type Error = D::Err; fn try_convtrans2d( self, @@ -169,7 +165,7 @@ where padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result { + ) -> Result { let (img, filters) = self; let (inp_chan, h, w) = img.shape; let img = img.try_reshape_like(&(Const::<1>, inp_chan, h, w))?; @@ -229,7 +225,6 @@ where D, T, >; - type Error = D::Err; fn try_convtrans2d( self, @@ -237,7 +232,7 @@ where padding: Padding, dilation: Dilation, groups: Groups, - ) -> Result { + ) -> Result { let (img, filters) = self; assert_eq!(img.shape.1, filters.shape.0); assert_eq!(filters.shape.2, filters.shape.3); diff --git a/dfdx-core/src/tensor_ops/cos/mod.rs b/dfdx-core/src/tensor_ops/cos/mod.rs index 01bd69d4c..a18be20b4 100644 --- a/dfdx-core/src/tensor_ops/cos/mod.rs +++ b/dfdx-core/src/tensor_ops/cos/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor Result { + pub fn try_cos(self) -> Result { try_unary_op(CosKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/div/mod.rs b/dfdx-core/src/tensor_ops/div/mod.rs index 1907b3971..41b0fe586 100644 --- a/dfdx-core/src/tensor_ops/div/mod.rs +++ b/dfdx-core/src/tensor_ops/div/mod.rs @@ -47,9 +47,9 @@ where } /// Fallible version of [std::ops::Div]. See [div] -pub trait TryDiv: HasErr { +pub trait TryDiv { type Output; - fn try_div(self, rhs: Rhs) -> Result; + fn try_div(self, rhs: Rhs) -> Result; } impl, R> TryDiv> @@ -60,7 +60,7 @@ where { type Output = Self; /// See [div] - fn try_div(self, rhs: Tensor) -> Result { + fn try_div(self, rhs: Tensor) -> Result { try_binary_op(BinaryDivKernelOp, self, rhs) } } @@ -71,7 +71,7 @@ where { type Output = Self; /// See [div] - fn try_div(self, rhs: Rhs) -> Result { + fn try_div(self, rhs: Rhs) -> Result { let rhs: f64 = rhs.into(); let scalar = E::from_f64(rhs).unwrap(); try_unary_op(ScalarDivKernelOp { scalar }, self) diff --git a/dfdx-core/src/tensor_ops/dropout/cpu_kernel.rs b/dfdx-core/src/tensor_ops/dropout/cpu_kernel.rs index 3216a1fa9..4eed47cab 100644 --- a/dfdx-core/src/tensor_ops/dropout/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/dropout/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Dtype, Shape}, - tensor::{unique_id, Cpu, Tensor}, + tensor::{unique_id, Cpu, Error, Tensor}, }; use num_traits::Float; @@ -12,7 +12,7 @@ impl super::DropoutKernel for Cpu { &self, op: super::DropoutKernelOp, inp: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut rng = StdRng::seed_from_u64(op.seed); let dist = Bernoulli::new(op.prob).unwrap(); let mut out = Tensor { @@ -39,7 +39,7 @@ impl super::DropoutKernel for Cpu { inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut rng = StdRng::seed_from_u64(op.seed); let dist = Bernoulli::new(op.prob).unwrap(); debug_assert_eq!(grad_inp.len(), grad_out.len()); diff --git a/dfdx-core/src/tensor_ops/dropout/cuda_kernel.rs b/dfdx-core/src/tensor_ops/dropout/cuda_kernel.rs index fe66e508e..c13125f1a 100644 --- a/dfdx-core/src/tensor_ops/dropout/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/dropout/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use std::vec::Vec; @@ -48,7 +48,7 @@ where &self, op: super::DropoutKernelOp, inp: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mask = { let mut rng = StdRng::seed_from_u64(op.seed); let dist = Bernoulli::new(op.prob).unwrap(); @@ -78,7 +78,7 @@ where inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mask = { let mut rng = StdRng::seed_from_u64(op.seed); let dist = Bernoulli::new(op.prob).unwrap(); diff --git a/dfdx-core/src/tensor_ops/dropout/mod.rs b/dfdx-core/src/tensor_ops/dropout/mod.rs index 1528853db..0ef0a10fb 100644 --- a/dfdx-core/src/tensor_ops/dropout/mod.rs +++ b/dfdx-core/src/tensor_ops/dropout/mod.rs @@ -3,10 +3,7 @@ mod cpu_kernel; #[cfg(feature = "cuda")] mod cuda_kernel; -use crate::{ - shapes::*, - tensor::{PutTape, RandomU64, SplitTape, Storage, Tape, Tensor}, -}; +use crate::{shapes::*, tensor::*}; #[repr(C)] #[derive(Debug, Clone, Copy)] @@ -20,14 +17,14 @@ pub trait DropoutKernel: Storage + RandomU64 { &self, op: DropoutKernelOp, inp: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, op: DropoutKernelOp, inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Zeros elements with probability `p` and scales all elements by `1 / (1 - p)`. @@ -62,7 +59,7 @@ impl, T: Tape> Tensor self.try_dropout(prob).unwrap() } /// See [dropout] - pub fn try_dropout(self, prob: impl Into) -> Result { + pub fn try_dropout(self, prob: impl Into) -> Result { let seed = self.device.random_u64(); let prob = prob.into(); let op = DropoutKernelOp { seed, prob }; diff --git a/dfdx-core/src/tensor_ops/exp/mod.rs b/dfdx-core/src/tensor_ops/exp/mod.rs index d17c56f72..3d04959d0 100644 --- a/dfdx-core/src/tensor_ops/exp/mod.rs +++ b/dfdx-core/src/tensor_ops/exp/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor Result { + pub fn try_exp(self) -> Result { try_unary_op(ExpKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/fast_gelu/mod.rs b/dfdx-core/src/tensor_ops/fast_gelu/mod.rs index 831ef5db6..6a4b46df5 100644 --- a/dfdx-core/src/tensor_ops/fast_gelu/mod.rs +++ b/dfdx-core/src/tensor_ops/fast_gelu/mod.rs @@ -48,7 +48,7 @@ impl, T: Tape> Ten self.try_fast_gelu().unwrap() } /// See [fast_gelu] - pub fn try_fast_gelu(self) -> Result { + pub fn try_fast_gelu(self) -> Result { try_unary_op(FastGeLUKernelOp, self) } @@ -60,7 +60,7 @@ impl, T: Tape> Ten /// Use [Tensor::try_fast_gelu] instead #[deprecated(since = "0.12.0", note = "Use `Tensor::try_fast_gelu` instead")] - pub fn try_gelu(self) -> Result { + pub fn try_gelu(self) -> Result { self.try_fast_gelu() } } diff --git a/dfdx-core/src/tensor_ops/huber_error/mod.rs b/dfdx-core/src/tensor_ops/huber_error/mod.rs index b2ee13c39..38bd27b93 100644 --- a/dfdx-core/src/tensor_ops/huber_error/mod.rs +++ b/dfdx-core/src/tensor_ops/huber_error/mod.rs @@ -50,7 +50,7 @@ impl, T: Tape> Tensor { self, rhs: Tensor, delta: impl Into, - ) -> Result + ) -> Result where T: Merge, { diff --git a/dfdx-core/src/tensor_ops/ln/mod.rs b/dfdx-core/src/tensor_ops/ln/mod.rs index a144e3b6f..2e1ae0673 100644 --- a/dfdx-core/src/tensor_ops/ln/mod.rs +++ b/dfdx-core/src/tensor_ops/ln/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor Result { + pub fn try_ln(self) -> Result { try_unary_op(LnKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/log_softmax.rs b/dfdx-core/src/tensor_ops/log_softmax.rs index 9a8023899..487c33e5a 100644 --- a/dfdx-core/src/tensor_ops/log_softmax.rs +++ b/dfdx-core/src/tensor_ops/log_softmax.rs @@ -38,7 +38,7 @@ impl, T: Tape> Tensor { self.try_log_softmax::().unwrap() } /// See [log_softmax()] - pub fn try_log_softmax(self) -> Result + pub fn try_log_softmax(self) -> Result where S: ReduceShape, { diff --git a/dfdx-core/src/tensor_ops/logsumexp_to.rs b/dfdx-core/src/tensor_ops/logsumexp_to.rs index 8f43abf4f..c8573638b 100644 --- a/dfdx-core/src/tensor_ops/logsumexp_to.rs +++ b/dfdx-core/src/tensor_ops/logsumexp_to.rs @@ -2,7 +2,7 @@ use super::*; use crate::{shapes::*, tensor::*}; /// Reduction along multiple axes using [LogSumExp](https://en.wikipedia.org/wiki/LogSumExp). -pub trait LogSumExpTo: HasErr + HasShape { +pub trait LogSumExpTo: Sized + HasShape { /// [LogSumExp](https://en.wikipedia.org/wiki/LogSumExp) reduction. /// /// **Pytorch equivalent**: `t.exp().sum(Axes).log()` @@ -31,13 +31,13 @@ pub trait LogSumExpTo: HasErr + HasShape { self.try_logsumexp().unwrap() } /// Fallible version of [LogSumExpTo::logsumexp] - fn try_logsumexp(self) -> Result, Self::Err> + fn try_logsumexp(self) -> Result, Error> where Self::Shape: ReduceShapeTo; } impl, T: Tape> LogSumExpTo for Tensor { - fn try_logsumexp(self) -> Result, Self::Err> + fn try_logsumexp(self) -> Result, Error> where Self::Shape: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/matmul/cpu_kernel.rs b/dfdx-core/src/tensor_ops/matmul/cpu_kernel.rs index 9e9b497ae..bf3e6ce0e 100644 --- a/dfdx-core/src/tensor_ops/matmul/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/matmul/cpu_kernel.rs @@ -1,7 +1,7 @@ #![allow(clippy::needless_return)] use crate::shapes::*; -use crate::tensor::{Cpu, Tensor, ZerosTensor}; +use crate::tensor::{Cpu, Error, Tensor, ZerosTensor}; use std::sync::Arc; @@ -236,7 +236,7 @@ where &self, lhs: &Tensor<(M, K), E, Self>, rhs: &Tensor<(K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let (m, k) = lhs.shape; let n = rhs.shape.1; let mut out = self.try_zeros_like(&(m, n))?; @@ -259,7 +259,7 @@ where rhs: &Tensor<(K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (m, k) = lhs.shape; let n = rhs.shape.1; let strides = (m, n).strides(); @@ -295,7 +295,7 @@ where &self, lhs: &Tensor<(B, M, K), E, Self>, rhs: &Tensor<(K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let (batch, m, k) = lhs.shape; let n = rhs.shape.1; let mut out = self.try_zeros_like(&(batch, m, n))?; @@ -321,7 +321,7 @@ where rhs: &Tensor<(K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (batch, m, k) = lhs.shape; let n = rhs.shape.1; let strides = (batch, m, n).strides(); @@ -359,7 +359,7 @@ where &self, lhs: &Tensor<(B, M, K), E, Self>, rhs: &Tensor<(B, K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let (b, m, k) = lhs.shape; let n = rhs.shape.2; let mut out = self.try_zeros_like(&(b, m, n))?; @@ -387,7 +387,7 @@ where rhs: &Tensor<(B, K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (b, m, k) = lhs.shape; let n = rhs.shape.2; let strides = (b, m, n).strides(); @@ -425,7 +425,7 @@ where &self, lhs: &Tensor<(B, S, M, K), E, Self>, rhs: &Tensor<(B, S, K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let (b, s, m, k) = lhs.shape; let n = rhs.shape.3; let mut out = self.try_zeros_like(&(b, s, m, n))?; @@ -453,7 +453,7 @@ where rhs: &Tensor<(B, S, K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (b, s, m, k) = lhs.shape; let n = rhs.shape.3; let strides = (b, s, m, n).strides(); diff --git a/dfdx-core/src/tensor_ops/matmul/cuda_kernel.rs b/dfdx-core/src/tensor_ops/matmul/cuda_kernel.rs index b6787d848..179ef06f2 100644 --- a/dfdx-core/src/tensor_ops/matmul/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/matmul/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{cuda::Cuda, Tensor}, + tensor::{cuda::Cuda, Error, Tensor}, }; use cudarc::{ @@ -245,7 +245,7 @@ where &self, lhs: &Tensor<(M, K), E, Self>, rhs: &Tensor<(K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let (m, _) = lhs.shape; let (k, n) = rhs.shape; let shape = (m, n); @@ -275,7 +275,7 @@ where rhs: &Tensor<(K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (m, _) = lhs.shape; let (k, n) = rhs.shape; let strides = (m, n).strides(); @@ -320,7 +320,7 @@ where &self, lhs: &Tensor<(B, M, K), E, Self>, rhs: &Tensor<(K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let (batch, m, _) = lhs.shape; let (k, n) = rhs.shape; let shape = (batch, m, n); @@ -347,7 +347,7 @@ where rhs: &Tensor<(K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (batch, m, _) = lhs.shape; let (k, n) = rhs.shape; let strides = (batch, m, n).strides(); @@ -396,7 +396,7 @@ where &self, lhs: &Tensor<(B, M, K), E, Self>, rhs: &Tensor<(B, K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { assert_ne!(lhs.strides[0], 0); assert_ne!(rhs.strides[0], 0); let (batch, m, _) = lhs.shape; @@ -425,7 +425,7 @@ where rhs: &Tensor<(B, K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (batch, m, _) = lhs.shape; let (_, k, n) = rhs.shape; let strides = (batch, m, n).strides(); @@ -470,7 +470,7 @@ where &self, lhs: &Tensor<(B, S, M, K), E, Self>, rhs: &Tensor<(B, S, K, N), E, Self>, - ) -> Result, Self::Err> { + ) -> Result, Error> { assert_ne!(lhs.strides[0], 0); assert_ne!(rhs.strides[0], 0); assert_ne!(lhs.strides[1], 0); @@ -527,7 +527,7 @@ where rhs: &Tensor<(B, S, K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let (batch, seq, m, _) = lhs.shape; let (_, _, k, n) = rhs.shape; let strides = (batch, seq, m, n).strides(); diff --git a/dfdx-core/src/tensor_ops/matmul/mod.rs b/dfdx-core/src/tensor_ops/matmul/mod.rs index 5247f4f88..b5ece54b2 100644 --- a/dfdx-core/src/tensor_ops/matmul/mod.rs +++ b/dfdx-core/src/tensor_ops/matmul/mod.rs @@ -7,7 +7,7 @@ pub(super) mod cuda_kernel; use crate::{ shapes::{Const, Dim, Dtype, Shape}, - tensor::{HasErr, Merge, PutTape, SplitTape, Storage, Tape, Tensor}, + tensor::{Error, Merge, PutTape, SplitTape, Storage, Tape, Tensor}, }; use super::reshape_to::{ReshapeKernel, ReshapeTo}; @@ -68,12 +68,12 @@ where } /// Fallible matrix multiplication. See [matmul] for examples. -pub trait TryMatMul: HasErr { +pub trait TryMatMul: Sized { type Output; fn matmul(self, rhs: Rhs) -> Self::Output { self.try_matmul(rhs).unwrap() } - fn try_matmul(self, rhs: Rhs) -> Result; + fn try_matmul(self, rhs: Rhs) -> Result; } #[rustfmt::skip] @@ -85,14 +85,14 @@ fn try_binary_op< D: Storage, RhsTape: Tape, LhsTape: Tape + Merge, - Fwd: 'static + FnMut(&D, &Tensor, &Tensor) -> Result, D::Err>, - Bwd: 'static + FnMut(&D, &Tensor, &mut D::Vec, &Tensor, &mut D::Vec, &D::Vec) -> Result<(), D::Err>, + Fwd: 'static + FnMut(&D, &Tensor, &Tensor) -> Result, crate::tensor::Error>, + Bwd: 'static + FnMut(&D, &Tensor, &mut D::Vec, &Tensor, &mut D::Vec, &D::Vec) -> Result<(), crate::tensor::Error>, >( lhs: Tensor, rhs: Tensor, mut fwd: Fwd, mut bwd: Bwd, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { let (lhs, ltape) = lhs.split_tape(); let (rhs, rtape) = rhs.split_tape(); let mut tape = ltape.merge(rtape); @@ -115,7 +115,7 @@ pub trait MatMatKernel: Storage { &self, lhs: &Tensor<(M, K), E, Self>, rhs: &Tensor<(K, N), E, Self>, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, @@ -124,7 +124,7 @@ pub trait MatMatKernel: Storage { rhs: &Tensor<(K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl + Merge, R: Tape> @@ -133,7 +133,7 @@ where D: MatMatKernel + ReshapeKernel, { type Output = Tensor<(M, N), E, D, T>; - fn try_matmul(self, rhs: Tensor<(N,), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(N,), E, D, R>) -> Result { let m = self.shape.0; let n = rhs.shape.0; let lhs = self.try_reshape_like(&(m, Const::<1>))?; @@ -147,7 +147,7 @@ where D: MatMatKernel + ReshapeKernel, { type Output = Tensor<(N,), E, D, T>; - fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result { let k1 = self.shape.0; let (k2, n) = rhs.shape; assert_eq!(k1, k2); @@ -162,7 +162,7 @@ where D: MatMatKernel + ReshapeKernel, { type Output = Tensor<(M,), E, D, T>; - fn try_matmul(self, rhs: Tensor<(K,), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(K,), E, D, R>) -> Result { let (m, k1) = self.shape; let k2 = rhs.shape.0; assert_eq!(k1, k2); @@ -185,7 +185,7 @@ where /// let y: Tensor, f32, _> = dev.zeros(); /// let _: Tensor, f32, _> = x.try_matmul(y); /// ``` - fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result { assert_eq!(self.shape.1, rhs.shape.0); try_binary_op(self, rhs, D::forward, D::backward) } @@ -196,7 +196,7 @@ pub trait MatMatBrKernel: Storage { &self, lhs: &Tensor<(B, M, K), E, Self>, rhs: &Tensor<(K, N), E, Self>, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, @@ -205,7 +205,7 @@ pub trait MatMatBrKernel: Storage { rhs: &Tensor<(K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl, T, R> @@ -222,7 +222,7 @@ where /// let y: Tensor, f32, _> = dev.zeros(); /// let _: Tensor, f32, _> = x.try_matmul(y); /// ``` - fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(K, N), E, D, R>) -> Result { assert_eq!(self.shape.2, rhs.shape.0); try_binary_op(self, rhs, D::forward, D::backward) } @@ -233,7 +233,7 @@ pub trait MatMatBatch3Kernel: Storage { &self, lhs: &Tensor<(B, M, K), E, Self>, rhs: &Tensor<(B, K, N), E, Self>, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, @@ -242,7 +242,7 @@ pub trait MatMatBatch3Kernel: Storage { rhs: &Tensor<(B, K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl TryMatMul> @@ -260,7 +260,7 @@ where /// let y: Tensor, f32, _> = dev.zeros(); /// let _: Tensor, f32, _> = x.try_matmul(y); /// ``` - fn try_matmul(self, rhs: Tensor<(B, K, N), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(B, K, N), E, D, R>) -> Result { assert_eq!(self.shape.0, rhs.shape.0); assert_eq!(self.shape.2, rhs.shape.1); try_binary_op(self, rhs, D::forward, D::backward) @@ -272,7 +272,7 @@ pub trait MatMatBatch4Kernel: Storage { &self, lhs: &Tensor<(B, S, M, K), E, Self>, rhs: &Tensor<(B, S, K, N), E, Self>, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, @@ -281,7 +281,7 @@ pub trait MatMatBatch4Kernel: Storage { rhs: &Tensor<(B, S, K, N), E, Self>, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl @@ -299,7 +299,7 @@ where /// let y: Tensor, f32, _> = dev.zeros(); /// let _: Tensor, f32, _> = x.try_matmul(y); /// ``` - fn try_matmul(self, rhs: Tensor<(B, S, K, N), E, D, R>) -> Result { + fn try_matmul(self, rhs: Tensor<(B, S, K, N), E, D, R>) -> Result { assert_eq!(self.shape.0, rhs.shape.0); assert_eq!(self.shape.1, rhs.shape.1); assert_eq!(self.shape.3, rhs.shape.2); diff --git a/dfdx-core/src/tensor_ops/max_to/cpu_kernel.rs b/dfdx-core/src/tensor_ops/max_to/cpu_kernel.rs index 049028bac..1308e84f6 100644 --- a/dfdx-core/src/tensor_ops/max_to/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/max_to/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Axes, Dtype, HasAxes, ReduceShapeTo, Shape}, - tensor::{Cpu, Tensor, ZerosTensor}, + tensor::{Cpu, Error, Tensor, ZerosTensor}, tensor_ops::utilities::reduction_utils::index_for_reductions, }; @@ -11,7 +11,7 @@ impl super::MaxReduceKernel for Cpu { &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo, { @@ -44,7 +44,7 @@ impl super::MaxReduceKernel for Cpu { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/max_to/cuda_kernel.rs b/dfdx-core/src/tensor_ops/max_to/cuda_kernel.rs index 31e384a6a..1926a24b3 100644 --- a/dfdx-core/src/tensor_ops/max_to/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/max_to/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, tensor_ops::reduction_utils::*, }; @@ -51,7 +51,7 @@ where &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo, { @@ -103,7 +103,7 @@ where grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/max_to/mod.rs b/dfdx-core/src/tensor_ops/max_to/mod.rs index 93eac94b3..462c6e561 100644 --- a/dfdx-core/src/tensor_ops/max_to/mod.rs +++ b/dfdx-core/src/tensor_ops/max_to/mod.rs @@ -10,7 +10,7 @@ pub trait MaxReduceKernel: Storage { &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo; fn backward( @@ -19,13 +19,13 @@ pub trait MaxReduceKernel: Storage { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo; } /// Reduction along multiple axes using `max`. -pub trait MaxTo: HasErr + HasShape { +pub trait MaxTo: Sized + HasShape { /// Max reduction. **Pytorch equivalent**: `t.amax(Ax)` /// /// **NOTE** This evenly distributes gradients between all equal maximum values, instead @@ -55,13 +55,13 @@ pub trait MaxTo: HasErr + HasShape { self.try_max().unwrap() } /// Fallible version of [MaxTo::max] - fn try_max(self) -> Result, Self::Err> + fn try_max(self) -> Result, Error> where Self::Shape: ReduceShapeTo; } impl, T: Tape> MaxTo for Tensor { - fn try_max(self) -> Result, Self::Err> + fn try_max(self) -> Result, Error> where Self::Shape: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/maximum/mod.rs b/dfdx-core/src/tensor_ops/maximum/mod.rs index 4d042c468..3ef180774 100644 --- a/dfdx-core/src/tensor_ops/maximum/mod.rs +++ b/dfdx-core/src/tensor_ops/maximum/mod.rs @@ -39,7 +39,10 @@ impl, LTape: Tape> Tensor } /// See [maximum] - pub fn try_maximum(self, rhs: Tensor) -> Result + pub fn try_maximum( + self, + rhs: Tensor, + ) -> Result where LTape: Merge, { diff --git a/dfdx-core/src/tensor_ops/mean_to.rs b/dfdx-core/src/tensor_ops/mean_to.rs index fb0b2f851..136d64d53 100644 --- a/dfdx-core/src/tensor_ops/mean_to.rs +++ b/dfdx-core/src/tensor_ops/mean_to.rs @@ -2,7 +2,7 @@ use super::*; use crate::{shapes::*, tensor::*}; /// Reduction along multiple axes using `mean`. -pub trait MeanTo: HasErr + HasShape { +pub trait MeanTo: Sized + HasShape { /// Mean reduction. **Pytorch equivalent**: `t.mean(Axes)` /// /// Example: @@ -29,13 +29,13 @@ pub trait MeanTo: HasErr + HasShape { self.try_mean().unwrap() } /// Fallible version of [MeanTo::mean] - fn try_mean(self) -> Result, Self::Err> + fn try_mean(self) -> Result, Error> where Self::Shape: HasAxes + ReduceShapeTo; } impl, T: Tape> MeanTo for Tensor { - fn try_mean(self) -> Result, Self::Err> + fn try_mean(self) -> Result, Error> where Self::Shape: HasAxes + ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/min_to/cpu_kernel.rs b/dfdx-core/src/tensor_ops/min_to/cpu_kernel.rs index e44302622..9b7f35491 100644 --- a/dfdx-core/src/tensor_ops/min_to/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/min_to/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Axes, Dtype, HasAxes, ReduceShapeTo, Shape}, - tensor::{Cpu, Tensor, ZerosTensor}, + tensor::{Cpu, Error, Tensor, ZerosTensor}, tensor_ops::utilities::reduction_utils::index_for_reductions, }; @@ -11,7 +11,7 @@ impl super::MinReduceKernel for Cpu { &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo, { @@ -44,7 +44,7 @@ impl super::MinReduceKernel for Cpu { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/min_to/cuda_kernel.rs b/dfdx-core/src/tensor_ops/min_to/cuda_kernel.rs index 8db3c1b5a..1e493b12e 100644 --- a/dfdx-core/src/tensor_ops/min_to/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/min_to/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, tensor_ops::reduction_utils::*, }; @@ -51,7 +51,7 @@ where &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo, { @@ -103,7 +103,7 @@ where grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/min_to/mod.rs b/dfdx-core/src/tensor_ops/min_to/mod.rs index ae426f4f1..9cef2da54 100644 --- a/dfdx-core/src/tensor_ops/min_to/mod.rs +++ b/dfdx-core/src/tensor_ops/min_to/mod.rs @@ -10,7 +10,7 @@ pub trait MinReduceKernel: Storage { &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo; fn backward( @@ -19,13 +19,13 @@ pub trait MinReduceKernel: Storage { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo; } /// Reduction along multiple axes using `min`. -pub trait MinTo: HasErr + HasShape { +pub trait MinTo: Sized + HasShape { /// Min reduction. **Pytorch equivalent**: `t.amin(Ax)` /// /// **NOTE** This evenly distributes gradients between all equal maximum values, instead @@ -55,13 +55,13 @@ pub trait MinTo: HasErr + HasShape { self.try_min().unwrap() } /// Fallible version of [MinTo::min] - fn try_min(self) -> Result, Self::Err> + fn try_min(self) -> Result, Error> where Self::Shape: ReduceShapeTo; } impl, T: Tape> MinTo for Tensor { - fn try_min(self) -> Result, Self::Err> + fn try_min(self) -> Result, Error> where Self::Shape: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/minimum/mod.rs b/dfdx-core/src/tensor_ops/minimum/mod.rs index e93c38426..adcc6dfbc 100644 --- a/dfdx-core/src/tensor_ops/minimum/mod.rs +++ b/dfdx-core/src/tensor_ops/minimum/mod.rs @@ -39,7 +39,10 @@ impl, LTape: Tape> Tensor } /// See [minimum] - pub fn try_minimum(self, rhs: Tensor) -> Result + pub fn try_minimum( + self, + rhs: Tensor, + ) -> Result where LTape: Merge, { diff --git a/dfdx-core/src/tensor_ops/mul/mod.rs b/dfdx-core/src/tensor_ops/mul/mod.rs index 0d345de14..8179509d3 100644 --- a/dfdx-core/src/tensor_ops/mul/mod.rs +++ b/dfdx-core/src/tensor_ops/mul/mod.rs @@ -46,9 +46,9 @@ where } /// Fallible version of [std::ops::Mul]. See [mul]. -pub trait TryMul: HasErr { +pub trait TryMul { type Output; - fn try_mul(self, rhs: Rhs) -> Result; + fn try_mul(self, rhs: Rhs) -> Result; } impl, LhsTape: Tape, R> @@ -57,7 +57,7 @@ where LhsTape: Merge, { type Output = Self; - fn try_mul(self, rhs: Tensor) -> Result { + fn try_mul(self, rhs: Tensor) -> Result { try_binary_op(BinaryMulKernelOp, self, rhs) } } @@ -67,7 +67,7 @@ where D: UnaryKernel, E>, { type Output = Self; - fn try_mul(self, rhs: Rhs) -> Result { + fn try_mul(self, rhs: Rhs) -> Result { let rhs: f64 = rhs.into(); let scalar: E = E::from_f64(rhs).unwrap(); try_unary_op(ScalarMulKernelOp { scalar }, self) diff --git a/dfdx-core/src/tensor_ops/nans_to/mod.rs b/dfdx-core/src/tensor_ops/nans_to/mod.rs index 9b511d5e7..176bb1c5e 100644 --- a/dfdx-core/src/tensor_ops/nans_to/mod.rs +++ b/dfdx-core/src/tensor_ops/nans_to/mod.rs @@ -35,7 +35,7 @@ impl, E>, T: Tape> Te self.try_nans_to(value).unwrap() } /// See [nans_to] - pub fn try_nans_to(self, value: impl Into) -> Result { + pub fn try_nans_to(self, value: impl Into) -> Result { let value = E::from_f64(value.into()).unwrap(); try_unary_op(NansToKernelOp(value), self) } diff --git a/dfdx-core/src/tensor_ops/negate/mod.rs b/dfdx-core/src/tensor_ops/negate/mod.rs index 89bf985c6..f6ad27db0 100644 --- a/dfdx-core/src/tensor_ops/negate/mod.rs +++ b/dfdx-core/src/tensor_ops/negate/mod.rs @@ -30,7 +30,7 @@ impl, T: Tape> Tenso pub fn negate(self) -> Self { self.try_negate().unwrap() } - pub fn try_negate(self) -> Result { + pub fn try_negate(self) -> Result { try_unary_op(NegateKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/normalize.rs b/dfdx-core/src/tensor_ops/normalize.rs index 49e13994d..46e2b0fd5 100644 --- a/dfdx-core/src/tensor_ops/normalize.rs +++ b/dfdx-core/src/tensor_ops/normalize.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Axes, Dtype, ReduceShape, Shape}, - tensor::{HasErr, Tape, Tensor}, + tensor::{Error, Tape, Tensor}, }; use super::{BroadcastTo, Device, MeanTo, TryAdd, TryDiv, TrySub}; @@ -32,10 +32,7 @@ impl, T: Tape> Tensor { } /// See [normalize] - pub fn try_normalize( - self, - epsilon: impl Into, - ) -> Result::Err> + pub fn try_normalize(self, epsilon: impl Into) -> Result where S: ReduceShape, { diff --git a/dfdx-core/src/tensor_ops/permute_to.rs b/dfdx-core/src/tensor_ops/permute_to.rs index 920d46e62..372ec9df2 100644 --- a/dfdx-core/src/tensor_ops/permute_to.rs +++ b/dfdx-core/src/tensor_ops/permute_to.rs @@ -21,7 +21,7 @@ use crate::{shapes::*, tensor::*}; /// let b: Tensor, f32, _> = a.permute::<_, Axes2<1, 0>>(); /// assert_eq!(b.array(), [[1.0, 4.0], [2.0, 5.0], [3.0, 6.0]]); /// ``` -pub trait PermuteTo: HasErr + HasShape { +pub trait PermuteTo: Sized + HasShape { /// Permutes the tensor. fn permute(self) -> Self::WithShape where @@ -30,13 +30,13 @@ pub trait PermuteTo: HasErr + HasShape { self.try_permute().unwrap() } /// Fallible version of [PermuteTo::permute] - fn try_permute(self) -> Result, Self::Err> + fn try_permute(self) -> Result, Error> where Self::Shape: PermuteShapeTo; } impl, T: Tape> PermuteTo for Tensor { - fn try_permute(self) -> Result, Self::Err> + fn try_permute(self) -> Result, Error> where Self::Shape: PermuteShapeTo, { diff --git a/dfdx-core/src/tensor_ops/pool2d/cpu_kernel.rs b/dfdx-core/src/tensor_ops/pool2d/cpu_kernel.rs index 478b19738..aaefa32eb 100644 --- a/dfdx-core/src/tensor_ops/pool2d/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/pool2d/cpu_kernel.rs @@ -59,7 +59,7 @@ impl super::Pool2DKind { } impl super::Pool2DKernel for Cpu { - fn alloc(&self, s: S) -> Result, Self::Err> { + fn alloc(&self, s: S) -> Result, Error> { self.try_zeros_like(&s) } fn forward( @@ -67,7 +67,7 @@ impl super::Pool2DKernel for Cpu { op: super::Pool2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let istr = make_4d::(inp.strides); let ostr = make_4d::(out.strides); @@ -106,7 +106,7 @@ impl super::Pool2DKernel for Cpu { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let istr = make_4d::(inp.strides); let ostr = make_4d::(out.strides); diff --git a/dfdx-core/src/tensor_ops/pool2d/cuda_kernel.rs b/dfdx-core/src/tensor_ops/pool2d/cuda_kernel.rs index 720a516a0..6ab840ab9 100644 --- a/dfdx-core/src/tensor_ops/pool2d/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/pool2d/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use std::sync::Arc; @@ -51,7 +51,7 @@ impl super::Pool2DKernel for Cuda where Self: HasCudaKernel, { - fn alloc(&self, s: S) -> Result, Self::Err> { + fn alloc(&self, s: S) -> Result, Error> { let data = unsafe { self.alloc_empty::(s.num_elements()) }?; Ok(self.build_tensor(s, s.strides(), data)) } @@ -60,7 +60,7 @@ where op: super::Pool2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::FWD, Self::FWD) { self.dev .load_ptx(PTX_SRC.into(), Self::FWD, &[Self::FWD, Self::BWD])?; @@ -87,7 +87,7 @@ where grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let inp_strides = self.dev.htod_copy(make_4d::(inp.strides).into())?; let out_strides = self.dev.htod_copy(make_4d::(out.strides).into())?; let bwd_fn = self.dev.get_func(Self::FWD, Self::BWD).unwrap(); diff --git a/dfdx-core/src/tensor_ops/pool2d/mod.rs b/dfdx-core/src/tensor_ops/pool2d/mod.rs index 0281d8971..150525c70 100644 --- a/dfdx-core/src/tensor_ops/pool2d/mod.rs +++ b/dfdx-core/src/tensor_ops/pool2d/mod.rs @@ -32,14 +32,14 @@ pub struct Pool2DOp { } pub(super) trait Pool2DKernel: Storage { - fn alloc(&self, s: S) -> Result, Self::Err>; + fn alloc(&self, s: S) -> Result, Error>; fn forward( &self, op: Pool2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; #[allow(clippy::too_many_arguments)] fn backward( @@ -49,12 +49,11 @@ pub(super) trait Pool2DKernel: Storage { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } pub trait TryPool2D: Sized { type Pooled; - type Error: std::fmt::Debug; fn pool2d( self, @@ -75,7 +74,7 @@ pub trait TryPool2D: Sized { stride: Stride, padding: Padding, dilation: Dilation, - ) -> Result; + ) -> Result; } impl< @@ -89,7 +88,6 @@ where Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>: Sized, { type Pooled = Const<{ (DIM + 2 * PADDING - DILATION * (KERNEL - 1) - 1) / STRIDE + 1 }>; - type Error = std::convert::Infallible; fn try_pool2d( self, _: Pool2DKind, @@ -97,7 +95,7 @@ where _: Const, _: Const, _: Const, - ) -> Result { + ) -> Result { Ok(Const) } } @@ -106,7 +104,6 @@ impl TryPool2D for usize { type Pooled = usize; - type Error = std::convert::Infallible; fn try_pool2d( self, _: Pool2DKind, @@ -114,7 +111,7 @@ impl stride: Stride, padding: Padding, dilation: Dilation, - ) -> Result { + ) -> Result { Ok((self + 2 * padding.size() - 1) .checked_sub(dilation.size() * (kernel.size() - 1)) .unwrap() @@ -140,7 +137,6 @@ where T: Tape, { type Pooled = Tensor<(Chan, H::Pooled, W::Pooled), E, D, T>; - type Error = D::Err; fn try_pool2d( self, @@ -149,7 +145,7 @@ where stride: Stride, padding: Padding, dilation: Dilation, - ) -> Result { + ) -> Result { let (chan, h, w) = self.shape; let img = self.try_reshape_like(&(Const::<1>, chan, h, w))?; let out = img.try_pool2d(kind, kernel, stride, padding, dilation)?; @@ -176,7 +172,6 @@ where T: Tape, { type Pooled = Tensor<(Batch, Chan, H::Pooled, W::Pooled), E, D, T>; - type Error = D::Err; fn try_pool2d( self, @@ -185,7 +180,7 @@ where stride: Stride, padding: Padding, dilation: Dilation, - ) -> Result { + ) -> Result { let (batch, chan, h, w) = self.shape; if self.strides != self.shape.strides() { panic!("Image input to pool2d must be contiguous"); diff --git a/dfdx-core/src/tensor_ops/pow/cuda_kernel.rs b/dfdx-core/src/tensor_ops/pow/cuda_kernel.rs index e698e2f4d..7d58027e3 100644 --- a/dfdx-core/src/tensor_ops/pow/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/pow/cuda_kernel.rs @@ -41,7 +41,7 @@ where &self, op: super::PowiKernelOp, inp: Cow>, - ) -> Result, Self::Err> { + ) -> Result, Error> { self.forward(super::PowfKernelOp(E::from_i32(op.0).unwrap()), inp) } @@ -52,7 +52,7 @@ where grad_inp: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { self.backward( super::PowfKernelOp(E::from_i32(op.0).unwrap()), inp, diff --git a/dfdx-core/src/tensor_ops/pow/mod.rs b/dfdx-core/src/tensor_ops/pow/mod.rs index 1d33ad243..bce8cff1e 100644 --- a/dfdx-core/src/tensor_ops/pow/mod.rs +++ b/dfdx-core/src/tensor_ops/pow/mod.rs @@ -34,7 +34,7 @@ impl, E>, T: Tape> Tens self.try_powf(exponent).unwrap() } /// See [powf] - pub fn try_powf(self, exponent: impl Into) -> Result { + pub fn try_powf(self, exponent: impl Into) -> Result { let exponent = E::from_f64(exponent.into()).unwrap(); try_unary_op(PowfKernelOp(exponent), self) } @@ -60,7 +60,7 @@ impl, T: Tape> Tensor< self.try_powi(exponent).unwrap() } /// See [powi] - pub fn try_powi(self, exponent: i32) -> Result { + pub fn try_powi(self, exponent: i32) -> Result { try_unary_op(PowiKernelOp(exponent), self) } } diff --git a/dfdx-core/src/tensor_ops/prelu.rs b/dfdx-core/src/tensor_ops/prelu.rs index 0bc018ca5..29d81c472 100644 --- a/dfdx-core/src/tensor_ops/prelu.rs +++ b/dfdx-core/src/tensor_ops/prelu.rs @@ -49,12 +49,12 @@ pub fn leakyrelu, T: Tape>( /// let r = prelu(t, a); /// assert_eq!(r.array(), [-0.05, 0.0, 1.0, 2.0]); /// ``` -pub trait TryPReLU: HasErr { +pub trait TryPReLU: Sized { fn prelu(self, rhs: T) -> Self { self.try_prelu(rhs).unwrap() } - fn try_prelu(self, rhs: T) -> Result; + fn try_prelu(self, rhs: T) -> Result; } impl, R> TryPReLU> @@ -64,7 +64,7 @@ where LhsTape: Merge, { /// See [prelu] - fn try_prelu(self, rhs: Tensor) -> Result { + fn try_prelu(self, rhs: Tensor) -> Result { let scaled = self.with_empty_tape().try_mul(rhs)?; self.try_lt(E::default())?.try_choose(scaled, self) } @@ -72,7 +72,7 @@ where impl, T: Tape> TryPReLU for Tensor { /// See [prelu] - fn try_prelu(self, rhs: E) -> Result { + fn try_prelu(self, rhs: E) -> Result { let dev = self.device.clone(); let scale = dev.tensor(rhs).retaped::().broadcast_like(self.shape()); let scaled = self.with_empty_tape().try_mul(scale)?; diff --git a/dfdx-core/src/tensor_ops/realize_to.rs b/dfdx-core/src/tensor_ops/realize_to.rs index 455ed62f1..894b2e86e 100644 --- a/dfdx-core/src/tensor_ops/realize_to.rs +++ b/dfdx-core/src/tensor_ops/realize_to.rs @@ -13,7 +13,7 @@ use crate::{shapes::*, tensor::*}; /// Err(old) => println!("Shape could not be realized, returned the original tensor"), /// } /// ``` -pub trait RealizeTo: HasErr + HasShape { +pub trait RealizeTo: Sized + HasShape { /// Realizes the concrete shape of the tensor as another compatable shape, /// or returns the original tensor if the new shape's dimensions are incompatable. fn realize::Shape as Shape>::Concrete>>( diff --git a/dfdx-core/src/tensor_ops/recip/mod.rs b/dfdx-core/src/tensor_ops/recip/mod.rs index 3a1d9157d..359226331 100644 --- a/dfdx-core/src/tensor_ops/recip/mod.rs +++ b/dfdx-core/src/tensor_ops/recip/mod.rs @@ -31,7 +31,7 @@ impl, T: Tape> Tensor self.try_recip().unwrap() } /// See [recip] - pub fn try_recip(self) -> Result { + pub fn try_recip(self) -> Result { try_unary_op(RecipKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/relu/mod.rs b/dfdx-core/src/tensor_ops/relu/mod.rs index 0701741c3..1b8dc8b1e 100644 --- a/dfdx-core/src/tensor_ops/relu/mod.rs +++ b/dfdx-core/src/tensor_ops/relu/mod.rs @@ -34,7 +34,7 @@ impl, T: Tape> Tensor< self.try_relu().unwrap() } /// See [relu] - pub fn try_relu(self) -> Result { + pub fn try_relu(self) -> Result { try_unary_op(ReLUKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/reshape_to/cpu_kernel.rs b/dfdx-core/src/tensor_ops/reshape_to/cpu_kernel.rs index 15a00e75d..31bad33a7 100644 --- a/dfdx-core/src/tensor_ops/reshape_to/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/reshape_to/cpu_kernel.rs @@ -1,7 +1,7 @@ use crate::shapes::{Dtype, Shape}; use crate::tensor::{ cpu::{LendingIterator, NdIndex}, - Cpu, Tensor, ZerosTensor, + Cpu, Error, Tensor, ZerosTensor, }; impl super::ReshapeKernel for Cpu { @@ -9,7 +9,7 @@ impl super::ReshapeKernel for Cpu { &self, dst: &Dst, inp: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut out = self.try_zeros_like(dst)?; let mut inp_iter = inp.iter(); let mut out_iter = out.iter_mut(); @@ -24,7 +24,7 @@ impl super::ReshapeKernel for Cpu { inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut inp_idx = NdIndex::new(inp.shape, inp.strides); let mut out_idx = NdIndex::new(*dst, dst.strides()); while let Some((i, o)) = inp_idx.next().zip(out_idx.next()) { diff --git a/dfdx-core/src/tensor_ops/reshape_to/cuda_kernel.rs b/dfdx-core/src/tensor_ops/reshape_to/cuda_kernel.rs index 607434f36..33612c01e 100644 --- a/dfdx-core/src/tensor_ops/reshape_to/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/reshape_to/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::{ driver::{DeviceSlice, LaunchAsync}, @@ -15,7 +15,7 @@ impl super::ReshapeKernel for Cuda { &self, dst: &Dst, inp: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let module = std::format!("reshape_fwd_{}", E::NAME); if !self.dev.has_func(&module, "reshape_fwd") { let src = FWD_KERNEL.replace("$T", E::NAME); @@ -62,7 +62,7 @@ impl super::ReshapeKernel for Cuda { inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let module = std::format!("reshape_bwd_{}", E::NAME); if !self.dev.has_func(&module, "reshape_bwd") { let src = BWD_KERNEL.replace("$T", E::NAME); diff --git a/dfdx-core/src/tensor_ops/reshape_to/mod.rs b/dfdx-core/src/tensor_ops/reshape_to/mod.rs index 04f058313..2d32ef460 100644 --- a/dfdx-core/src/tensor_ops/reshape_to/mod.rs +++ b/dfdx-core/src/tensor_ops/reshape_to/mod.rs @@ -10,14 +10,14 @@ pub trait ReshapeKernel: Storage { &self, dst: &Dst, inp: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, dst: &Dst, inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Changes the shape of a tensor without re-ordering axes. If the tensor is contiguous @@ -47,7 +47,7 @@ pub trait ReshapeKernel: Storage { /// let t: Tensor, f32, _> = dev.zeros(); /// let t: Tensor<(usize, ), f32, _> = t.reshape_like(&(8, )); /// ``` -pub trait ReshapeTo: HasErr + HasShape { +pub trait ReshapeTo: Sized + HasShape { /// Reshapes a tensor to a different compile time shape. fn reshape(self) -> Self::WithShape where @@ -57,7 +57,7 @@ pub trait ReshapeTo: HasErr + HasShape { self.try_reshape().unwrap() } /// Reshapes a tensor to a different compile time shape. - fn try_reshape(self) -> Result, Self::Err> + fn try_reshape(self) -> Result, Error> where Self::Shape: ConstShape, { @@ -75,16 +75,16 @@ pub trait ReshapeTo: HasErr + HasShape { self.try_contiguous().unwrap() } /// See [`ReshapeTo::contiguous`] - fn try_contiguous(self) -> Result, Self::Err> { + fn try_contiguous(self) -> Result, Error> { let shape = *self.shape(); self.try_reshape_like(&shape) } /// Reshapes a tensor to a different runtime shape. - fn try_reshape_like(self, dst: &Dst) -> Result, Self::Err>; + fn try_reshape_like(self, dst: &Dst) -> Result, Error>; } impl, T: Tape> ReshapeTo for Tensor { - fn try_reshape_like(self, dst: &Dst) -> Result, Self::Err> { + fn try_reshape_like(self, dst: &Dst) -> Result, Error> { assert_eq!(self.shape().num_elements(), dst.num_elements()); if self.shape.strides() == self.strides { Ok(Tensor { diff --git a/dfdx-core/src/tensor_ops/rmsprop/cpu_kernel.rs b/dfdx-core/src/tensor_ops/rmsprop/cpu_kernel.rs index f7408651c..d6c7f70fd 100644 --- a/dfdx-core/src/tensor_ops/rmsprop/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/rmsprop/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ dtypes::{Dtype, NotMixedPrecision}, - tensor::cpu::Cpu, + tensor::{cpu::Cpu, Error}, }; use super::{RMSpropConfig, RMSpropKernel, WeightDecay}; @@ -15,7 +15,7 @@ impl RMSpropKernel> for Cpu { square_avg: &mut Self::Vec, grad_avg: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let alpha = cfg.alpha as f32; let eps = cfg.eps as f32; let lr = cfg.lr as f32; @@ -80,7 +80,7 @@ impl RMSpropKernel for Cpu square_avg: &mut Self::Vec, grad_avg: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let alpha = E::from_f64(cfg.alpha).unwrap(); let eps = E::from_f64(cfg.eps).unwrap(); let lr = E::from_f64(cfg.lr).unwrap(); diff --git a/dfdx-core/src/tensor_ops/rmsprop/cuda_kernel.rs b/dfdx-core/src/tensor_ops/rmsprop/cuda_kernel.rs index 8491d44c3..ca1842e70 100644 --- a/dfdx-core/src/tensor_ops/rmsprop/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/rmsprop/cuda_kernel.rs @@ -1,7 +1,7 @@ use super::RMSpropConfig; use crate::{ dtypes::*, - tensor::{launch_cfg, Cuda}, + tensor::{launch_cfg, Cuda, Error}, tensor_ops::optim::*, }; @@ -82,7 +82,7 @@ where square_avg: &mut Self::Vec, grad_avg: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FWD) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, &[Self::FWD])?; } diff --git a/dfdx-core/src/tensor_ops/rmsprop/mod.rs b/dfdx-core/src/tensor_ops/rmsprop/mod.rs index 55afb3089..0ccddad39 100644 --- a/dfdx-core/src/tensor_ops/rmsprop/mod.rs +++ b/dfdx-core/src/tensor_ops/rmsprop/mod.rs @@ -55,7 +55,7 @@ pub trait RMSpropKernel: Storage { square_avg: &mut Self::Vec, grad_avg: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl RMSpropConfig { @@ -67,7 +67,7 @@ impl RMSpropConfig { square_avg: &mut D::Vec, grad_avg: &mut D::Vec, grad: &D::Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { param.device.rmsprop_kernel( self, std::sync::Arc::make_mut(&mut param.data), diff --git a/dfdx-core/src/tensor_ops/roll/cpu_kernel.rs b/dfdx-core/src/tensor_ops/roll/cpu_kernel.rs index f1814c701..f77e7d058 100644 --- a/dfdx-core/src/tensor_ops/roll/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/roll/cpu_kernel.rs @@ -10,7 +10,7 @@ impl super::RollKernel for Cpu { &self, op: super::RollOp, inp: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { let dims = inp.shape.concrete(); let strides = inp.shape.strides(); let mut data = self.try_alloc_zeros::(inp.shape.num_elements())?; @@ -39,7 +39,7 @@ impl super::RollKernel for Cpu { inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let dims = inp.shape.concrete(); let strides = inp.shape.strides(); let mut idx = NdIndex::new(inp.shape, inp.strides); diff --git a/dfdx-core/src/tensor_ops/roll/cuda_kernel.rs b/dfdx-core/src/tensor_ops/roll/cuda_kernel.rs index aaf3bcf6c..4404ceb7f 100644 --- a/dfdx-core/src/tensor_ops/roll/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/roll/cuda_kernel.rs @@ -32,7 +32,7 @@ where &self, op: super::RollOp, inp: &Tensor, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(Self::FNS[0], Self::FNS[0]) { self.dev.load_ptx(PTX_SRC.into(), Self::FNS[0], Self::FNS)?; } @@ -66,7 +66,7 @@ where inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let numel = inp.shape.num_elements(); let strides = inp.shape.strides(); diff --git a/dfdx-core/src/tensor_ops/roll/mod.rs b/dfdx-core/src/tensor_ops/roll/mod.rs index ad7318792..40a3ac5ec 100644 --- a/dfdx-core/src/tensor_ops/roll/mod.rs +++ b/dfdx-core/src/tensor_ops/roll/mod.rs @@ -19,14 +19,14 @@ pub trait RollKernel: Storage { &self, op: RollOp, inp: &Tensor, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, op: RollOp, inp: &Tensor, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Shifts data along an axis by a specified amount. @@ -47,7 +47,7 @@ pub trait RollKernel: Storage { /// let r = t.roll::>(1); /// assert_eq!(r.array(), [4.0, 1.0, 2.0, 3.0]); /// ``` -pub trait Roll: HasShape + HasErr { +pub trait Roll: Sized + HasShape { /// Shifts data along an axis by a specified amount. fn roll>(self, amount: usize) -> Self where @@ -57,13 +57,16 @@ pub trait Roll: HasShape + HasErr { } /// Shifts data along an axis by a specified amount. - fn try_roll>(self, amount: usize) -> Result + fn try_roll>(self, amount: usize) -> Result where Self::Shape: HasAxes; } impl, T: Tape> Roll for Tensor { - fn try_roll>(self, amount: usize) -> Result + fn try_roll>( + self, + amount: usize, + ) -> Result where S: HasAxes, { diff --git a/dfdx-core/src/tensor_ops/select_and_gather/cpu_kernel.rs b/dfdx-core/src/tensor_ops/select_and_gather/cpu_kernel.rs index 5b7d0ec7d..6acf39105 100644 --- a/dfdx-core/src/tensor_ops/select_and_gather/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/select_and_gather/cpu_kernel.rs @@ -3,7 +3,7 @@ use crate::shapes::{Axes, Dtype, RemoveDimTo, ReplaceDimTo, Shape}; use crate::tensor::{ cpu::{index_to_i, LendingIterator, NdIndex}, - Cpu, Storage, Tensor, ZerosTensor, + Cpu, Error, Storage, Tensor, ZerosTensor, }; impl super::ReplaceDimKernel for Cpu { @@ -11,7 +11,7 @@ impl super::ReplaceDimKernel for Cpu { &self, inp: &Tensor, idx: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReplaceDimTo, { @@ -53,7 +53,7 @@ impl super::ReplaceDimKernel for Cpu { idx: &Tensor, out: &Tensor, grad_out: &>::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReplaceDimTo, { @@ -87,7 +87,7 @@ impl super::RemoveDimKernel for Cpu { &self, inp: &Tensor, idx: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: RemoveDimTo, { @@ -126,7 +126,7 @@ impl super::RemoveDimKernel for Cpu { idx: &Tensor, out: &Tensor, grad_out: &>::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: RemoveDimTo, { diff --git a/dfdx-core/src/tensor_ops/select_and_gather/cuda_kernel.rs b/dfdx-core/src/tensor_ops/select_and_gather/cuda_kernel.rs index 1d4ce5b5b..a7a5133ae 100644 --- a/dfdx-core/src/tensor_ops/select_and_gather/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/select_and_gather/cuda_kernel.rs @@ -2,7 +2,7 @@ use crate::{ dtypes::*, shapes::{RemoveDimTo, ReplaceDimTo, Shape}, - tensor::{launch_cfg, Cuda, Storage, Tensor}, + tensor::{launch_cfg, Cuda, Error, Storage, Tensor}, }; use cudarc::driver::{DeviceSlice, LaunchAsync}; @@ -16,7 +16,7 @@ macro_rules! impl_cuda_kernels { &self, inp: &Tensor, idx: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReplaceDimTo, { @@ -65,7 +65,7 @@ macro_rules! impl_cuda_kernels { idx: &Tensor, _: &Tensor, grad_out: &>::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReplaceDimTo, { @@ -101,7 +101,7 @@ macro_rules! impl_cuda_kernels { &self, inp: &Tensor, idx: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: RemoveDimTo, { @@ -153,7 +153,7 @@ macro_rules! impl_cuda_kernels { idx: &Tensor, out: &Tensor, grad_out: &>::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: RemoveDimTo, { diff --git a/dfdx-core/src/tensor_ops/select_and_gather/mod.rs b/dfdx-core/src/tensor_ops/select_and_gather/mod.rs index f6178e9c7..2a56f1acf 100644 --- a/dfdx-core/src/tensor_ops/select_and_gather/mod.rs +++ b/dfdx-core/src/tensor_ops/select_and_gather/mod.rs @@ -12,7 +12,7 @@ pub trait ReplaceDimKernel: Storage + Storage { &self, inp: &Tensor, idx: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReplaceDimTo; fn backward( @@ -22,7 +22,7 @@ pub trait ReplaceDimKernel: Storage + Storage { idx: &Tensor, out: &Tensor, grad_out: &>::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReplaceDimTo; } @@ -32,7 +32,7 @@ pub trait RemoveDimKernel: Storage + Storage { &self, inp: &Tensor, idx: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: RemoveDimTo; fn backward( @@ -42,7 +42,7 @@ pub trait RemoveDimKernel: Storage + Storage { idx: &Tensor, out: &Tensor, grad_out: &>::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: RemoveDimTo; } @@ -73,7 +73,7 @@ pub trait RemoveDimKernel: Storage + Storage { /// let idx: Tensor, usize, _> = dev.tensor([0, 2, 4]); /// let _: Tensor, f32, _> = a.select(idx); ///``` -pub trait SelectTo + Storage>: HasErr + HasShape { +pub trait SelectTo + Storage>: Sized + HasShape { /// Select values given indices. fn select(self, idx: Tensor) -> Self::WithShape where @@ -86,7 +86,7 @@ pub trait SelectTo + Storage>: HasErr + HasShape { fn try_select( self, idx: Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: RemoveDimTo; } @@ -97,7 +97,7 @@ impl, T: Tape> SelectTo fn try_select( self, idx: Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: RemoveDimTo, { @@ -146,7 +146,7 @@ impl, T: Tape> SelectTo /// let idx: Tensor, usize, _> = dev.tensor([[0, 1], [2, 3], [4, 4]]); /// let _: Tensor, f32, _> = a.gather(idx); ///``` -pub trait GatherTo + Storage>: HasErr + HasShape { +pub trait GatherTo + Storage>: Sized + HasShape { /// Gather values given indices. fn gather(self, idx: Tensor) -> Self::WithShape where @@ -158,7 +158,7 @@ pub trait GatherTo + Storage>: HasErr + HasShape { fn try_gather( self, idx: Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: ReplaceDimTo; } @@ -169,7 +169,7 @@ impl, T: Tape> GatherTo fn try_gather( self, idx: Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: ReplaceDimTo, { diff --git a/dfdx-core/src/tensor_ops/sgd/cpu_kernel.rs b/dfdx-core/src/tensor_ops/sgd/cpu_kernel.rs index 69c5653c4..2fcae414b 100644 --- a/dfdx-core/src/tensor_ops/sgd/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/sgd/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ dtypes::{Dtype, NotMixedPrecision}, - tensor::cpu::*, + tensor::{cpu::*, Error}, }; use super::{Momentum, SgdConfig, SgdKernel, WeightDecay}; @@ -13,7 +13,7 @@ impl SgdKernel> for Cpu { param: &mut Self::Vec, velocity: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let lr = cfg.lr as f32; for ((p, g), v) in param @@ -62,7 +62,7 @@ impl SgdKernel for Cpu { param: &mut Self::Vec, velocity: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let lr = E::from_f64(cfg.lr).unwrap(); for ((p, mut g), v) in param diff --git a/dfdx-core/src/tensor_ops/sgd/cuda_kernel.rs b/dfdx-core/src/tensor_ops/sgd/cuda_kernel.rs index 6d29812fc..d893a325d 100644 --- a/dfdx-core/src/tensor_ops/sgd/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/sgd/cuda_kernel.rs @@ -2,7 +2,7 @@ use super::SgdConfig; use crate::{ dtypes::*, - tensor::{launch_cfg, Cuda}, + tensor::{launch_cfg, Cuda, Error}, tensor_ops::optim::*, }; @@ -71,7 +71,7 @@ where param: &mut Self::Vec, velocity: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FWD) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, &[Self::FWD])?; } diff --git a/dfdx-core/src/tensor_ops/sgd/mod.rs b/dfdx-core/src/tensor_ops/sgd/mod.rs index 3515f9f57..16737b2c5 100644 --- a/dfdx-core/src/tensor_ops/sgd/mod.rs +++ b/dfdx-core/src/tensor_ops/sgd/mod.rs @@ -5,7 +5,7 @@ mod cuda_kernel; use crate::{ shapes::{Dtype, Shape}, - tensor::{Storage, Tensor}, + tensor::{Error, Storage, Tensor}, }; use super::optim::{Momentum, WeightDecay}; @@ -90,7 +90,7 @@ pub trait SgdKernel: Storage { param: &mut Self::Vec, velocity: &mut Self::Vec, grad: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } impl SgdConfig { @@ -100,7 +100,7 @@ impl SgdConfig { param: &mut Tensor, velocity: &mut D::Vec, grad: &D::Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { param.device.sgd_kernel( self, std::sync::Arc::make_mut(&mut param.data), diff --git a/dfdx-core/src/tensor_ops/sigmoid/mod.rs b/dfdx-core/src/tensor_ops/sigmoid/mod.rs index 316b85564..ff41f64f5 100644 --- a/dfdx-core/src/tensor_ops/sigmoid/mod.rs +++ b/dfdx-core/src/tensor_ops/sigmoid/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tens self.try_sigmoid().unwrap() } /// See [sigmoid] - pub fn try_sigmoid(self) -> Result { + pub fn try_sigmoid(self) -> Result { try_unary_op(SigmoidKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/sin/mod.rs b/dfdx-core/src/tensor_ops/sin/mod.rs index 841d1ebba..035da953c 100644 --- a/dfdx-core/src/tensor_ops/sin/mod.rs +++ b/dfdx-core/src/tensor_ops/sin/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor Result { + pub fn try_sin(self) -> Result { try_unary_op(SinKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/slice/cpu_kernel.rs b/dfdx-core/src/tensor_ops/slice/cpu_kernel.rs index c987112cf..25d27ed76 100644 --- a/dfdx-core/src/tensor_ops/slice/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/slice/cpu_kernel.rs @@ -7,7 +7,7 @@ impl SliceKernel for Cpu { &self, inp: &Tensor, slice: &Slice, - ) -> Result, Self::Err> { + ) -> Result, Error> { let dst = inp.shape.slice(slice).unwrap(); let mut out = self.try_zeros_like(&dst)?; @@ -31,7 +31,7 @@ impl SliceKernel for Cpu { grad_inp: &mut Self::Vec, grad_out: &Self::Vec, slice: &Slice, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let dst = inp.shape.slice(slice).unwrap(); let mut inp_idx = NdIndex::new(dst, inp.strides); diff --git a/dfdx-core/src/tensor_ops/slice/cuda_kernel.rs b/dfdx-core/src/tensor_ops/slice/cuda_kernel.rs index 7e2f85d1a..d7548daaa 100644 --- a/dfdx-core/src/tensor_ops/slice/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/slice/cuda_kernel.rs @@ -2,7 +2,7 @@ use crate::{ dtypes::*, prelude::cpu::NdIndex, shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::driver::{CudaSlice, LaunchAsync}; @@ -56,7 +56,7 @@ where &self, inp: &Tensor, slice: &Slice, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(Self::MOD, Self::FNS[0]) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?; } @@ -94,7 +94,7 @@ where grad_inp: &mut Self::Vec, grad_out: &Self::Vec, slice: &Slice, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::MOD, Self::FNS[1]) { self.dev.load_ptx(PTX_SRC.into(), Self::MOD, Self::FNS)?; } diff --git a/dfdx-core/src/tensor_ops/slice/mod.rs b/dfdx-core/src/tensor_ops/slice/mod.rs index 9dce18815..db3a6854d 100644 --- a/dfdx-core/src/tensor_ops/slice/mod.rs +++ b/dfdx-core/src/tensor_ops/slice/mod.rs @@ -9,7 +9,7 @@ pub trait SliceKernel: Storage { &self, inp: &Tensor, slice: &Slice, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward, Slice>( &self, @@ -17,7 +17,7 @@ pub trait SliceKernel: Storage { grad_inp: &mut Self::Vec, grad_out: &Self::Vec, slice: &Slice, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } /// Slices all dimensions of a tensor, with the starting and ending indices of each dimension @@ -53,7 +53,10 @@ pub fn slice, E: Unit, D: SliceKernel, T: Tape, Sl impl, T: Tape> Tensor { /// Fallible version of [Tensor::slice] - pub fn try_slice(self, slice: Slice) -> Result, D::Err> + pub fn try_slice( + self, + slice: Slice, + ) -> Result, crate::tensor::Error> where S: SliceShape, Slice: 'static, diff --git a/dfdx-core/src/tensor_ops/softmax.rs b/dfdx-core/src/tensor_ops/softmax.rs index 4390319da..0a6ec8aab 100644 --- a/dfdx-core/src/tensor_ops/softmax.rs +++ b/dfdx-core/src/tensor_ops/softmax.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor { self.try_softmax::().unwrap() } /// See [softmax()] - pub fn try_softmax(self) -> Result + pub fn try_softmax(self) -> Result where S: ReduceShape, { diff --git a/dfdx-core/src/tensor_ops/sqrt/mod.rs b/dfdx-core/src/tensor_ops/sqrt/mod.rs index 6fa472f1c..b4703946e 100644 --- a/dfdx-core/src/tensor_ops/sqrt/mod.rs +++ b/dfdx-core/src/tensor_ops/sqrt/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor< self.try_sqrt().unwrap() } /// See [sqrt] - pub fn try_sqrt(self) -> Result { + pub fn try_sqrt(self) -> Result { try_unary_op(SqrtKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/square/mod.rs b/dfdx-core/src/tensor_ops/square/mod.rs index aadab4f17..ae9ae0da3 100644 --- a/dfdx-core/src/tensor_ops/square/mod.rs +++ b/dfdx-core/src/tensor_ops/square/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tenso self.try_square().unwrap() } /// See [square] - pub fn try_square(self) -> Result { + pub fn try_square(self) -> Result { try_unary_op(SquareKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/stack/cpu_kernel.rs b/dfdx-core/src/tensor_ops/stack/cpu_kernel.rs index 49a0174e9..3f09fa923 100644 --- a/dfdx-core/src/tensor_ops/stack/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/stack/cpu_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::*, - tensor::{unique_id, Cpu, Tensor}, + tensor::{unique_id, Cpu, Error, Tensor}, }; use std::vec::Vec; @@ -10,7 +10,7 @@ impl super::StackKernel for Cpu { &self, num: Num, inp: &[Tensor], - ) -> Result, Self::Err> + ) -> Result, Error> where S: super::AddDim, { @@ -52,7 +52,7 @@ impl super::StackKernel for Cpu { &self, mut grad_inp: Vec<&mut Self::Vec>, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let mut offset = 0; for item in grad_inp.drain(..) { for gi in item.iter_mut() { diff --git a/dfdx-core/src/tensor_ops/stack/cuda_kernel.rs b/dfdx-core/src/tensor_ops/stack/cuda_kernel.rs index d6a783fbc..827a2be17 100644 --- a/dfdx-core/src/tensor_ops/stack/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/stack/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::{ driver::{DeviceSlice, LaunchAsync}, @@ -14,7 +14,7 @@ impl super::StackKernel for Cuda { &self, num: Num, inps: &[Tensor], - ) -> Result, Self::Err> + ) -> Result, Error> where S: super::AddDim, { @@ -54,7 +54,7 @@ impl super::StackKernel for Cuda { &self, mut grad_inp: Vec<&mut Self::Vec>, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let module_name = std::format!("stack_bwd_{}", E::NAME); if !self.dev.has_func(&module_name, "stack_bwd") { let src = BWD_KERNEL.replace("$Ty", E::NAME); diff --git a/dfdx-core/src/tensor_ops/stack/mod.rs b/dfdx-core/src/tensor_ops/stack/mod.rs index 3a6f8350e..f8e13d590 100644 --- a/dfdx-core/src/tensor_ops/stack/mod.rs +++ b/dfdx-core/src/tensor_ops/stack/mod.rs @@ -32,14 +32,13 @@ mod cuda_kernel; /// ``` pub trait TryStack: Sized { type Stacked; - type Err: std::fmt::Debug; /// Stack an array or vec of tensors together along a new dimension. fn stack(self) -> Self::Stacked { self.try_stack().unwrap() } /// Fallible version of [TryStack::stack] - fn try_stack(self) -> Result; + fn try_stack(self) -> Result; } impl, T, const N: usize> TryStack for [Tensor; N] @@ -48,8 +47,7 @@ where T: Tape, { type Stacked = Tensor; - type Err = D::Err; - fn try_stack(self) -> Result { + fn try_stack(self) -> Result { try_stack(self) } } @@ -60,16 +58,14 @@ where T: Tape, { type Stacked = Tensor; - type Err = D::Err; - fn try_stack(self) -> Result { + fn try_stack(self) -> Result { try_stack(self) } } -impl> TryStack for (A, B) { +impl TryStack for (A, B) { type Stacked = (A::Stacked, B::Stacked); - type Err = A::Err; - fn try_stack(self) -> Result { + fn try_stack(self) -> Result { Ok((self.0.try_stack()?, self.1.try_stack()?)) } } @@ -115,19 +111,15 @@ pub trait StackKernel: Storage { &self, num: Num, inp: &[Tensor], - ) -> Result, Self::Err> + ) -> Result, Error> where S: AddDim; - fn backward( - &self, - grad_inp: Vec<&mut Self::Vec>, - grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + fn backward(&self, grad_inp: Vec<&mut Self::Vec>, grad_out: &Self::Vec) -> Result<(), Error>; } fn try_stack, T, Items>( items: Items, -) -> Result, D::Err> +) -> Result, crate::tensor::Error> where Items: Array>, S: AddDim, diff --git a/dfdx-core/src/tensor_ops/stddev_to.rs b/dfdx-core/src/tensor_ops/stddev_to.rs index 9e116cabb..07beaecd2 100644 --- a/dfdx-core/src/tensor_ops/stddev_to.rs +++ b/dfdx-core/src/tensor_ops/stddev_to.rs @@ -2,7 +2,7 @@ use super::*; use crate::{shapes::*, tensor::*}; /// Reduction along multiple axes using standard deviation. -pub trait StddevTo: HasErr + HasShape { +pub trait StddevTo: Sized + HasShape { /// Standard deviation reduction. /// /// **Pytorch equivalent**: `t.std(Axes, unbiased=False)` @@ -25,7 +25,7 @@ pub trait StddevTo: HasErr + HasShape { fn try_stddev( self, epsilon: impl Into, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: HasAxes + ReduceShapeTo; } @@ -34,7 +34,7 @@ impl, T: Tape> StddevTo for Tensor( self, epsilon: impl Into, - ) -> Result, Self::Err> + ) -> Result, Error> where Self::Shape: HasAxes + ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/sub/mod.rs b/dfdx-core/src/tensor_ops/sub/mod.rs index 1fdd35813..a7b82759d 100644 --- a/dfdx-core/src/tensor_ops/sub/mod.rs +++ b/dfdx-core/src/tensor_ops/sub/mod.rs @@ -47,9 +47,9 @@ where } /// Fallible version of [std::ops::Sub]. See [sub] -pub trait TrySub: HasErr { +pub trait TrySub { type Output; - fn try_sub(self, rhs: Rhs) -> Result; + fn try_sub(self, rhs: Rhs) -> Result; } impl, LTape: Tape, R> @@ -58,7 +58,7 @@ where LTape: Merge, { type Output = Self; - fn try_sub(self, rhs: Tensor) -> Result { + fn try_sub(self, rhs: Tensor) -> Result { try_binary_op(BinarySubKernelOp, self, rhs) } } @@ -68,7 +68,7 @@ where D: UnaryKernel, E>, { type Output = Self; - fn try_sub(self, rhs: Rhs) -> Result { + fn try_sub(self, rhs: Rhs) -> Result { let rhs: f64 = rhs.into(); let scalar = E::from_f64(rhs).unwrap(); try_unary_op(ScalarSubKernelOp { scalar }, self) diff --git a/dfdx-core/src/tensor_ops/sum_to/cpu_kernel.rs b/dfdx-core/src/tensor_ops/sum_to/cpu_kernel.rs index 5daa7cced..373622dcf 100644 --- a/dfdx-core/src/tensor_ops/sum_to/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/sum_to/cpu_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::{Dtype, NotMixedPrecision}, shapes::{Axes, HasAxes, ReduceShapeTo, Shape}, - tensor::{Cpu, Tensor, Tensorlike, ZerosTensor}, + tensor::{Cpu, Error, Tensor, Tensorlike, ZerosTensor}, tensor_ops::utilities::reduction_utils::index_for_reductions, }; @@ -11,7 +11,7 @@ impl super::SumKernel> for Cpu { &self, dst: Dst, inp: &Tensor, Self>, - ) -> Result, Self>, Self::Err> + ) -> Result, Self>, Error> where Src: ReduceShapeTo, { @@ -46,7 +46,7 @@ impl super::SumKernel> for Cpu { inp: &impl Tensorlike, Self>, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { @@ -75,7 +75,7 @@ impl super::SumKernel for Cpu { &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo, { @@ -108,7 +108,7 @@ impl super::SumKernel for Cpu { inp: &impl Tensorlike, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/sum_to/cuda_kernel.rs b/dfdx-core/src/tensor_ops/sum_to/cuda_kernel.rs index 05d1dec32..868909411 100644 --- a/dfdx-core/src/tensor_ops/sum_to/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/sum_to/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor, Tensorlike}, + tensor::{launch_cfg, Cuda, Error, Tensor, Tensorlike}, tensor_ops::reduction_utils::*, }; @@ -46,7 +46,7 @@ where &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo, { @@ -99,7 +99,7 @@ where inp: &impl Tensorlike, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/sum_to/mod.rs b/dfdx-core/src/tensor_ops/sum_to/mod.rs index d0dae445b..e99c37101 100644 --- a/dfdx-core/src/tensor_ops/sum_to/mod.rs +++ b/dfdx-core/src/tensor_ops/sum_to/mod.rs @@ -10,7 +10,7 @@ pub trait SumKernel: Storage { &self, dst: Dst, inp: &Tensor, - ) -> Result, Self::Err> + ) -> Result, Error> where Src: ReduceShapeTo; fn backward( @@ -19,13 +19,13 @@ pub trait SumKernel: Storage { inp: &impl Tensorlike, grad_inp: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> + ) -> Result<(), Error> where Src: ReduceShapeTo; } /// Reduction along multiple axes using `sum`. -pub trait SumTo: HasErr + HasShape { +pub trait SumTo: Sized + HasShape { /// Sum reduction. **Pytorch equivalent**: `t.sum(Ax)` /// /// Example reducing a single axis: @@ -52,13 +52,13 @@ pub trait SumTo: HasErr + HasShape { self.try_sum().unwrap() } /// Fallible version of [SumTo::sum] - fn try_sum(self) -> Result, Self::Err> + fn try_sum(self) -> Result, Error> where Self::Shape: ReduceShapeTo; } impl, T: Tape> SumTo for Tensor { - fn try_sum(self) -> Result, Self::Err> + fn try_sum(self) -> Result, Error> where Self::Shape: ReduceShapeTo, { diff --git a/dfdx-core/src/tensor_ops/tanh/mod.rs b/dfdx-core/src/tensor_ops/tanh/mod.rs index 7ddf191ba..2c5b96064 100644 --- a/dfdx-core/src/tensor_ops/tanh/mod.rs +++ b/dfdx-core/src/tensor_ops/tanh/mod.rs @@ -33,7 +33,7 @@ impl, T: Tape> Tensor< self.try_tanh().unwrap() } /// See [tanh] - pub fn try_tanh(self) -> Result { + pub fn try_tanh(self) -> Result { try_unary_op(TanhKernelOp, self) } } diff --git a/dfdx-core/src/tensor_ops/to_dtype/cpu_kernel.rs b/dfdx-core/src/tensor_ops/to_dtype/cpu_kernel.rs index a0f948052..e06e7e01e 100644 --- a/dfdx-core/src/tensor_ops/to_dtype/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/to_dtype/cpu_kernel.rs @@ -1,10 +1,10 @@ use num_traits::AsPrimitive; use std::{sync::Arc, vec::Vec}; -use crate::prelude::{cpu::CachableVec, Cpu, Shape, Tensor, Unit}; +use crate::prelude::{cpu::CachableVec, Cpu, Error, Shape, Tensor, Unit}; impl, E2: Unit> super::ToDtypeKernel for Cpu { - fn forward(inp: Tensor) -> Result, Self::Err> { + fn forward(inp: Tensor) -> Result, Error> { let data: &[E1] = inp.data.as_ref(); let data: Vec = data.iter().map(|x| (*x).as_()).collect(); let data = CachableVec { diff --git a/dfdx-core/src/tensor_ops/to_dtype/cuda_kernel.rs b/dfdx-core/src/tensor_ops/to_dtype/cuda_kernel.rs index 5e6d232c6..ed147b0da 100644 --- a/dfdx-core/src/tensor_ops/to_dtype/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/to_dtype/cuda_kernel.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Shape, Unit}, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use cudarc::{ driver::{DeviceSlice, LaunchAsync}, @@ -22,7 +22,7 @@ extern \"C\" __global__ void kernel(const size_t n, const $Src *inp, $Dst *out) }"; impl super::ToDtypeKernel for Cuda { - fn forward(inp: Tensor) -> Result, Self::Err> { + fn forward(inp: Tensor) -> Result, Error> { let module = std::format!("convert_{}_to_{}", E1::NAME, E2::NAME); let cuda = &inp.device; diff --git a/dfdx-core/src/tensor_ops/to_dtype/mod.rs b/dfdx-core/src/tensor_ops/to_dtype/mod.rs index 85b0b5ab0..5f0b7c999 100644 --- a/dfdx-core/src/tensor_ops/to_dtype/mod.rs +++ b/dfdx-core/src/tensor_ops/to_dtype/mod.rs @@ -2,10 +2,10 @@ mod cpu_kernel; #[cfg(feature = "cuda")] mod cuda_kernel; -use crate::prelude::{Shape, Storage, Tensor, Unit}; +use crate::prelude::{Error, Shape, Storage, Tensor, Unit}; pub trait ToDtypeKernel: Storage + Storage { - fn forward(inp: Tensor) -> Result, Self::Err>; + fn forward(inp: Tensor) -> Result, Error>; } /// Copies the elements of a tensor, converting its data to a different dtype. @@ -29,7 +29,7 @@ pub fn to_dtype>( } impl> Tensor { - pub fn try_to_dtype(self) -> Result, D::Err> + pub fn try_to_dtype(self) -> Result, crate::tensor::Error> where D: ToDtypeKernel, { diff --git a/dfdx-core/src/tensor_ops/tri.rs b/dfdx-core/src/tensor_ops/tri.rs index 3a03d1da9..c1f8fa451 100644 --- a/dfdx-core/src/tensor_ops/tri.rs +++ b/dfdx-core/src/tensor_ops/tri.rs @@ -1,5 +1,5 @@ use crate::shapes::{Dtype, Shape}; -use crate::tensor::{HasErr, Tape, Tensor, TriangleTensor}; +use crate::tensor::{Error, Tape, Tensor, TriangleTensor}; use super::TryMul; @@ -11,7 +11,7 @@ pub fn lower_tri, T: Tape>( diagonal: impl Into>, ) -> Tensor where - Tensor: TryMul, Output = Tensor> + HasErr, + Tensor: TryMul, Output = Tensor>, { t.lower_tri(diagonal) } @@ -24,20 +24,17 @@ pub fn upper_tri, T: Tape>( diagonal: impl Into>, ) -> Tensor where - Tensor: TryMul, Output = Tensor> + HasErr, + Tensor: TryMul, Output = Tensor>, { t.upper_tri(diagonal) } impl, T: Tape> Tensor where - Self: TryMul, Output = Self> + HasErr, + Self: TryMul, Output = Self>, { /// See [lower_tri] - pub fn try_lower_tri( - self, - diagonal: impl Into>, - ) -> Result::Err> { + pub fn try_lower_tri(self, diagonal: impl Into>) -> Result { let out = self .device .try_lower_tri_like(&self.shape, E::ONE, diagonal)?; @@ -50,10 +47,7 @@ where } /// See [upper_tri] - pub fn try_upper_tri( - self, - diagonal: impl Into>, - ) -> Result::Err> { + pub fn try_upper_tri(self, diagonal: impl Into>) -> Result { let out = self .device .try_upper_tri_like(&self.shape, E::ONE, diagonal)?; diff --git a/dfdx-core/src/tensor_ops/upscale2d/cpu_kernel.rs b/dfdx-core/src/tensor_ops/upscale2d/cpu_kernel.rs index 82083e513..efbff95ef 100644 --- a/dfdx-core/src/tensor_ops/upscale2d/cpu_kernel.rs +++ b/dfdx-core/src/tensor_ops/upscale2d/cpu_kernel.rs @@ -1,5 +1,5 @@ use crate::shapes::*; -use crate::tensor::{Cpu, Tensor}; +use crate::tensor::{Cpu, Error, Tensor}; use std::sync::Arc; @@ -23,7 +23,7 @@ impl op: super::Upscale2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let istr = make_4d::(inp.strides); let ostr = make_4d::(out.strides); @@ -56,7 +56,7 @@ impl grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let istr = make_4d::(inp.strides); let ostr = make_4d::(out.strides); @@ -87,7 +87,7 @@ impl super::Upscale2DKernel for Cpu { op: super::Upscale2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let istr = make_4d::(inp.strides); let ostr = make_4d::(out.strides); @@ -138,7 +138,7 @@ impl super::Upscale2DKernel for Cpu { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let istr = make_4d::(inp.strides); let ostr = make_4d::(out.strides); diff --git a/dfdx-core/src/tensor_ops/upscale2d/cuda_kernel.rs b/dfdx-core/src/tensor_ops/upscale2d/cuda_kernel.rs index 038a86ae8..9c98646bd 100644 --- a/dfdx-core/src/tensor_ops/upscale2d/cuda_kernel.rs +++ b/dfdx-core/src/tensor_ops/upscale2d/cuda_kernel.rs @@ -1,7 +1,7 @@ use crate::{ dtypes::*, shapes::*, - tensor::{launch_cfg, Cuda, Tensor}, + tensor::{launch_cfg, Cuda, Error, Tensor}, }; use std::sync::Arc; @@ -71,7 +71,7 @@ where op: super::Upscale2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { if !self.dev.has_func(Self::FWD, Self::FWD) { self.dev .load_ptx(PTX_SRC.into(), Self::FWD, &[Self::FWD, Self::BWD])?; @@ -96,7 +96,7 @@ where grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let strides = self.dev.htod_copy(make_4d::(inp.strides).into())?; let bwd_fn = self.dev.get_func(Self::FWD, Self::BWD).unwrap(); let cfg = launch_cfg::<128>(out.shape().num_elements() as u32); diff --git a/dfdx-core/src/tensor_ops/upscale2d/mod.rs b/dfdx-core/src/tensor_ops/upscale2d/mod.rs index 5cd849ddc..d0adce479 100644 --- a/dfdx-core/src/tensor_ops/upscale2d/mod.rs +++ b/dfdx-core/src/tensor_ops/upscale2d/mod.rs @@ -5,7 +5,7 @@ mod cuda_kernel; use crate::{ shapes::*, - tensor::{HasErr, PutTape, SplitTape, Storage, Tape, Tensor, ZerosTensor}, + tensor::{Error, PutTape, SplitTape, Storage, Tape, Tensor, ZerosTensor}, }; #[repr(C)] @@ -65,7 +65,7 @@ pub trait Upscale2DKernel: Storage { op: Upscale2DOp, inp: &Tensor, out: &mut Tensor, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; fn backward( &self, @@ -74,17 +74,17 @@ pub trait Upscale2DKernel: Storage { grad_inp: &mut Self::Vec, out: &Tensor, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } -pub trait GenericUpscale2D: HasErr { +pub trait GenericUpscale2D { type Output; fn generic_upscale2d_like( self, method: M, height: OH, width: OW, - ) -> Result, Self::Err>; + ) -> Result, Error>; } /// Upscales an image to a new shape. Valid methods of upscaling are: @@ -107,7 +107,7 @@ pub trait GenericUpscale2D: HasErr { /// let t: Tensor, f32, _> = dev.zeros(); /// let y: Tensor<(Const<3>, usize, usize), f32, _> = t.upscale2d_like(NearestNeighbor, 64, 64); /// ``` -pub trait TryUpscale2D { +pub trait TryUpscale2D: Sized { /// Upscale to compile time known dimensions. fn upscale2d( self, @@ -122,7 +122,7 @@ pub trait TryUpscale2D { fn try_upscale2d( self, method: M, - ) -> Result<>::Output, Const>, Self::Err> + ) -> Result<>::Output, Const>, Error> where Self: GenericUpscale2D, { @@ -146,7 +146,7 @@ pub trait TryUpscale2D { method: M, height: OH, width: OW, - ) -> Result<>::Output, Self::Err> + ) -> Result<>::Output, Error> where Self: GenericUpscale2D, { @@ -172,7 +172,7 @@ impl< _method: M, out_height: OH, out_width: OW, - ) -> Result, Self::Err> { + ) -> Result, Error> { let in_height = self.shape.1; let in_width = self.shape.2; @@ -216,7 +216,7 @@ impl< _method: M, out_height: OH, out_width: OW, - ) -> Result, Self::Err> { + ) -> Result, Error> { let in_height = self.shape.2; let in_width = self.shape.3; diff --git a/dfdx-core/src/tensor_ops/utilities/backward.rs b/dfdx-core/src/tensor_ops/utilities/backward.rs index 7dc01ccf7..f78981468 100644 --- a/dfdx-core/src/tensor_ops/utilities/backward.rs +++ b/dfdx-core/src/tensor_ops/utilities/backward.rs @@ -4,19 +4,19 @@ use crate::tensor::*; /// Runs backprop algorithm with all operations contained in the tape that `t` has. /// /// This function takes ownership of `self` and returns [Gradients]. -pub trait Backward>: HasErr { +pub trait Backward>: Sized { /// Runs backprop fn backward(self) -> Gradients { self.try_backward().unwrap() } /// Fallible version of [Backward::backward] - fn try_backward(self) -> Result, Self::Err>; + fn try_backward(self) -> Result, Error>; } impl> Backward for Tensor> { - fn try_backward(self) -> Result, Self::Err> { + fn try_backward(self) -> Result, Error> { let (t, mut tape) = self.split_tape(); let t_ghost = t.ghost(); tape.add_backward_op(move |grads| { @@ -33,7 +33,7 @@ impl> Backward impl> Backward for Tensor>>> { - fn try_backward(self) -> Result, Self::Err> { + fn try_backward(self) -> Result, Error> { let (t, tape) = self.split_tape(); let t_ghost = t.ghost(); let mut tape = tape.lock().unwrap(); diff --git a/dfdx-core/src/tensor_ops/utilities/cpu_kernels.rs b/dfdx-core/src/tensor_ops/utilities/cpu_kernels.rs index 56718a72c..0848c1edc 100644 --- a/dfdx-core/src/tensor_ops/utilities/cpu_kernels.rs +++ b/dfdx-core/src/tensor_ops/utilities/cpu_kernels.rs @@ -5,7 +5,7 @@ use crate::{ shapes::{Dtype, Shape}, tensor::{ cpu::{Cpu, LendingIterator, NdIndex}, - unique_id, Tensor, Tensorlike, ZerosTensor, + unique_id, Error, Tensor, Tensorlike, ZerosTensor, }, }; @@ -51,7 +51,7 @@ impl> UnaryKernel for Cpu { &self, op: Op, inp: Cow>, - ) -> Result, Self::Err> { + ) -> Result, Error> { let mut out = match inp { Cow::Borrowed(inp) => { // allocate a new data buffer @@ -84,7 +84,7 @@ impl> UnaryKernel for Cpu { grad_inp: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { match (inp.data(), out.data()) { (None, None) => { let df = op.const_df(); @@ -115,7 +115,7 @@ impl> BinaryKernel for Cpu { op: Op, lhs: Cow>, rhs: Cow>, - ) -> Result, Self::Err> { + ) -> Result, Error> { match (lhs, rhs) { (Cow::Borrowed(lhs), Cow::Borrowed(rhs)) => { let mut out = self.try_zeros_like(&lhs.shape)?; @@ -169,7 +169,7 @@ impl> BinaryKernel for Cpu { rhs: &impl Tensorlike, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { match (lhs.data(), rhs.data()) { (Some(lhs_buf), Some(rhs_buf)) => { let mut lhs_idx = NdIndex::new(*lhs.shape(), lhs.strides()); diff --git a/dfdx-core/src/tensor_ops/utilities/cuda_kernels.rs b/dfdx-core/src/tensor_ops/utilities/cuda_kernels.rs index 9ce4df1d6..b472b0253 100644 --- a/dfdx-core/src/tensor_ops/utilities/cuda_kernels.rs +++ b/dfdx-core/src/tensor_ops/utilities/cuda_kernels.rs @@ -67,7 +67,7 @@ impl + DeviceRepr> UnaryKernel for Cuda &self, op: K, inp: Cow>, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(K::MODULE_NAME, K::FWD_FN_NAME) { self.dev .load_ptx(K::PTX_SRC.into(), K::MODULE_NAME, &K::ALL_FN_NAMES)?; @@ -103,7 +103,7 @@ impl + DeviceRepr> UnaryKernel for Cuda grad_inp: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let bwd_fn = self.dev.get_func(K::MODULE_NAME, K::BWD_FN_NAME).unwrap(); match (inp.data(), out.data()) { (None, None) => { @@ -219,7 +219,7 @@ impl + DeviceRepr + Clone> BinaryKernel op: K, lhs: Cow>, rhs: Cow>, - ) -> Result, Self::Err> { + ) -> Result, Error> { if !self.dev.has_func(K::MODULE_NAME, K::FWD_FN_NAME) { self.dev .load_ptx(K::PTX_SRC.into(), K::MODULE_NAME, &K::ALL_FN_NAMES)?; @@ -326,7 +326,7 @@ impl + DeviceRepr + Clone> BinaryKernel rhs: &impl Tensorlike, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err> { + ) -> Result<(), Error> { let bwd_lhs_fn = self .dev .get_func(K::MODULE_NAME, K::BWD_LHS_FN_NAME) diff --git a/dfdx-core/src/tensor_ops/utilities/ops.rs b/dfdx-core/src/tensor_ops/utilities/ops.rs index 15a8250d6..e8a3dcfef 100644 --- a/dfdx-core/src/tensor_ops/utilities/ops.rs +++ b/dfdx-core/src/tensor_ops/utilities/ops.rs @@ -1,6 +1,6 @@ use crate::{ shapes::{Dtype, HasShape, Shape}, - tensor::{Merge, PutTape, SplitTape, Storage, Tape, Tensor, Tensorlike}, + tensor::*, }; use std::borrow::Cow; @@ -11,7 +11,7 @@ pub trait UnaryKernel: Storage { &self, op: Op, inp: Cow>, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, op: Op, @@ -19,7 +19,7 @@ pub trait UnaryKernel: Storage { grad_inp: &mut Self::Vec, out: &impl Tensorlike, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } pub trait BinaryKernel: Storage { @@ -29,7 +29,7 @@ pub trait BinaryKernel: Storage { op: Op, lhs: Cow>, rhs: Cow>, - ) -> Result, Self::Err>; + ) -> Result, Error>; fn backward( &self, op: Op, @@ -38,7 +38,7 @@ pub trait BinaryKernel: Storage { rhs: &impl Tensorlike, grad_rhs: &mut Self::Vec, grad_out: &Self::Vec, - ) -> Result<(), Self::Err>; + ) -> Result<(), Error>; } pub(crate) fn try_unary_op< @@ -50,7 +50,7 @@ pub(crate) fn try_unary_op< >( op: Op, inp: Tensor, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { let (inp, mut tape) = inp.split_tape(); let inp_ghost = inp.ghost(); let dev = inp.device.clone(); @@ -99,7 +99,7 @@ pub(crate) fn try_binary_op< op: Op, lhs: Tensor, rhs: Tensor, -) -> Result, D::Err> { +) -> Result, crate::tensor::Error> { assert_eq!(lhs.shape(), rhs.shape()); let (lhs, ltape) = lhs.split_tape(); let (rhs, rtape) = rhs.split_tape(); diff --git a/dfdx-core/src/tensor_ops/var_to.rs b/dfdx-core/src/tensor_ops/var_to.rs index 1a593655c..0a7ef618b 100644 --- a/dfdx-core/src/tensor_ops/var_to.rs +++ b/dfdx-core/src/tensor_ops/var_to.rs @@ -2,7 +2,7 @@ use super::*; use crate::{shapes::*, tensor::*}; /// Reduction alogn multiple axes using variance -pub trait VarTo: HasErr + HasShape { +pub trait VarTo: Sized + HasShape { /// Result [Tensor] has smaller number of dimensions. /// /// **Pytorch equivalent**: `t.var(Axes, unbiased=False)` @@ -22,13 +22,13 @@ pub trait VarTo: HasErr + HasShape { self.try_var().unwrap() } /// Fallible version of [VarTo::var] - fn try_var(self) -> Result, Self::Err> + fn try_var(self) -> Result, Error> where Self::Shape: HasAxes + ReduceShapeTo; } impl, T: Tape> VarTo for Tensor { - fn try_var(self) -> Result, Self::Err> + fn try_var(self) -> Result, Error> where Self::Shape: HasAxes + ReduceShapeTo, { diff --git a/dfdx-derives/src/lib.rs b/dfdx-derives/src/lib.rs index f258c679c..60da4982e 100644 --- a/dfdx-derives/src/lib.rs +++ b/dfdx-derives/src/lib.rs @@ -31,8 +31,7 @@ macro_rules! has_attr { /// pub struct Abs; /// impl, T: Tape> Module> for Abs { /// type Output = Tensor; -/// type Error = D::Err; -/// fn try_forward(&self, x: Tensor) -> Result { +/// fn try_forward(&self, x: Tensor) -> Result { /// x.try_abs() /// } /// } @@ -50,8 +49,7 @@ macro_rules! has_attr { /// for Reshape /// { /// type Output = Tensor; -/// type Error = D::Err; -/// fn try_forward(&self, x: Tensor) -> Result { +/// fn try_forward(&self, x: Tensor) -> Result { /// x.try_reshape_like(&self.0) /// } /// } @@ -77,8 +75,7 @@ macro_rules! has_attr { /// X: TryAdd, /// { /// type Output = X; -/// type Error = D::Err; -/// fn try_forward(&self, x: X) -> Result { +/// fn try_forward(&self, x: X) -> Result { /// self.matmul.try_forward(x.with_empty_tape())?.try_add(x) /// } /// } @@ -206,7 +203,7 @@ pub fn custom_module(input: proc_macro::TokenStream) -> proc_macro::TokenStream } impl #build_impl dfdx_core::nn_traits::ResetParams for #builder_name #built_ty #built_where { - fn try_reset_params(&mut self) -> Result<(), Dev::Err> { + fn try_reset_params(&mut self) -> Result<(), dfdx_core::tensor::Error> { Ok(()) } } @@ -217,13 +214,13 @@ pub fn custom_module(input: proc_macro::TokenStream) -> proc_macro::TokenStream optimizer: &mut Optim, gradients: &dfdx_core::tensor::Gradients, missing_tensors: &mut Vec, - ) -> Result<(), Dev::Err> { + ) -> Result<(), dfdx_core::tensor::Error> { Ok(()) } } impl #build_impl dfdx_core::nn_traits::ZeroGrads for #builder_name #built_ty #built_where { - fn try_zero_grads(&self, grads: &mut dfdx_core::tensor::Gradients) -> Result<(), Dev::Err> { + fn try_zero_grads(&self, grads: &mut dfdx_core::tensor::Gradients) -> Result<(), dfdx_core::tensor::Error> { Ok(()) } } @@ -261,7 +258,7 @@ pub fn custom_module(input: proc_macro::TokenStream) -> proc_macro::TokenStream quote! { impl #build_impl dfdx_core::nn_traits::BuildOnDevice for #builder_name #builder_ty #built_where { type Built = #built_name #built_ty; - fn try_build_on_device(&self, device: &Dev) -> Result { + fn try_build_on_device(&self, device: &Dev) -> Result { let built = #built_name { #(#recurse)* }; Ok(built) } @@ -280,7 +277,7 @@ pub fn custom_module(input: proc_macro::TokenStream) -> proc_macro::TokenStream quote! { impl #build_impl dfdx_core::nn_traits::BuildOnDevice for #builder_name #builder_ty #built_where { type Built = #built_name #built_ty; - fn try_build_on_device(&self, device: &Dev) -> Result { + fn try_build_on_device(&self, device: &Dev) -> Result { let built = #built_name(#(#recurse)*); Ok(built) } @@ -291,7 +288,7 @@ pub fn custom_module(input: proc_macro::TokenStream) -> proc_macro::TokenStream quote! { impl #build_impl dfdx_core::nn_traits::BuildOnDevice for #builder_name #builder_ty #built_where { type Built = #built_name #built_ty; - fn try_build_on_device(&self, device: &Dev) -> Result { + fn try_build_on_device(&self, device: &Dev) -> Result { Ok(#built_name) } } @@ -422,7 +419,7 @@ pub fn sequential(input: proc_macro::TokenStream) -> proc_macro::TokenStream { quote! { impl #built_impl dfdx_core::nn_traits::BuildOnDevice for #builder_name #builder_ty #built_where { type Built = #built_name #built_ty; - fn try_build_on_device(&self, device: &Dev) -> Result { + fn try_build_on_device(&self, device: &Dev) -> Result { let built = #built_name { #(#recurse)* }; @@ -439,7 +436,7 @@ pub fn sequential(input: proc_macro::TokenStream) -> proc_macro::TokenStream { quote! { impl #built_impl dfdx_core::nn_traits::BuildOnDevice for #builder_name #builder_ty #built_where { type Built = #built_name #built_ty; - fn try_build_on_device(&self, device: &Dev) -> Result { + fn try_build_on_device(&self, device: &Dev) -> Result { #built_name( #(#recurse)* ) @@ -455,12 +452,8 @@ pub fn sequential(input: proc_macro::TokenStream) -> proc_macro::TokenStream { // Get's the output type of the sequential. Also adds Module bounds to the where clause. let mut last_ty = quote!(Input); - let err = quote!(::Err); let output_ty = { let where_clause = module_generics.make_where_clause(); - where_clause - .predicates - .push(parse_quote!(Input: dfdx_core::prelude::HasErr)); match &input.data { Data::Struct(ref obj) => match obj.fields { Fields::Named(ref fields) => { @@ -471,7 +464,7 @@ pub fn sequential(input: proc_macro::TokenStream) -> proc_macro::TokenStream { .push(parse_quote!(#ty: dfdx_core::nn_traits::BuildOnDevice)); where_clause .predicates - .push(parse_quote!(<#ty as dfdx_core::nn_traits::BuildOnDevice>::Built: dfdx_core::nn_traits::Module<#last_ty, Error = #err>)); + .push(parse_quote!(<#ty as dfdx_core::nn_traits::BuildOnDevice>::Built: dfdx_core::nn_traits::Module<#last_ty>)); last_ty = parse_quote!(<<#ty as dfdx_core::nn_traits::BuildOnDevice>::Built as dfdx_core::nn_traits::Module<#last_ty>>::Output); }); } @@ -483,7 +476,7 @@ pub fn sequential(input: proc_macro::TokenStream) -> proc_macro::TokenStream { .push(parse_quote!(#ty: dfdx_core::nn_traits::BuildOnDevice)); where_clause .predicates - .push(parse_quote!(<#ty as dfdx_core::nn_traits::BuildOnDevice>::Built: dfdx_core::nn_traits::Module<#last_ty, Error = #err>)); + .push(parse_quote!(<#ty as dfdx_core::nn_traits::BuildOnDevice>::Built: dfdx_core::nn_traits::Module<#last_ty>)); last_ty = parse_quote!(<<#ty as dfdx_core::nn_traits::BuildOnDevice>::Built as dfdx_core::nn_traits::Module<#last_ty>>::Output); }); } @@ -523,8 +516,7 @@ pub fn sequential(input: proc_macro::TokenStream) -> proc_macro::TokenStream { quote! { impl #module_impl dfdx_core::nn_traits::Module for #built_name #built_ty #module_where { type Output = #output_ty; - type Error = #err; - fn try_forward(&self, x: Input) -> Result { + fn try_forward(&self, x: Input) -> Result { #src Ok(x) } @@ -606,7 +598,7 @@ pub fn reset_params(input: proc_macro::TokenStream) -> proc_macro::TokenStream { proc_macro::TokenStream::from(quote! { impl #impl_generics dfdx_core::nn_traits::ResetParams for #name #ty_generics #where_clause { - fn try_reset_params(&mut self) -> Result<(), Dev::Err> { + fn try_reset_params(&mut self) -> Result<(), dfdx_core::tensor::Error> { #resets Ok(()) } @@ -690,7 +682,7 @@ pub fn update_params(input: proc_macro::TokenStream) -> proc_macro::TokenStream optimizer: &mut Optim, gradients: &dfdx_core::tensor::Gradients, missing_tensors: &mut Vec, - ) -> Result<(), Dev::Err> { + ) -> Result<(), dfdx_core::tensor::Error> { #updates Ok(()) } @@ -773,7 +765,7 @@ pub fn zero_grads(input: proc_macro::TokenStream) -> proc_macro::TokenStream { proc_macro::TokenStream::from(quote! { impl #impl_generics dfdx_core::nn_traits::ZeroGrads for #name #ty_generics #where_clause { - fn try_zero_grads(&self, grads: &mut dfdx_core::prelude::Gradients) -> Result<(), Dev::Err> { + fn try_zero_grads(&self, grads: &mut dfdx_core::prelude::Gradients) -> Result<(), dfdx_core::tensor::Error> { #zero_grads Ok(()) } diff --git a/dfdx/examples/advanced-train-loop.rs b/dfdx/examples/advanced-train-loop.rs index 9e130a3fb..5c5ed0ed2 100644 --- a/dfdx/examples/advanced-train-loop.rs +++ b/dfdx/examples/advanced-train-loop.rs @@ -12,7 +12,7 @@ fn classification_train< Lbl, // Our model just needs to implement these two things! ModuleMut for forward // and TensorCollection for optimizer/alloc_grads/zero_grads - Model: Module + ZeroGrads + UpdateParams, + Model: Module + ZeroGrads + UpdateParams, // optimizer, pretty straight forward Opt: Optimizer, // our data will just be any iterator over these items. easy! @@ -22,7 +22,7 @@ fn classification_train< Criterion: FnMut(Model::Output, Lbl) -> Loss, // the Loss needs to be able to call backward, and we also use // this generic as an output - Loss: Backward + AsArray, + Loss: Backward + AsArray, // Dtype & Device to tie everything together E: Dtype, D: Device, @@ -32,7 +32,7 @@ fn classification_train< mut criterion: Criterion, data: Data, batch_accum: usize, -) -> Result<(), D::Err> { +) -> Result<(), Error> { let mut grads = model.try_alloc_grads()?; for (i, (inp, lbl)) in data.enumerate() { let y = model.try_forward_mut(inp.traced(grads))?; diff --git a/dfdx/src/nn/layers/abs.rs b/dfdx/src/nn/layers/abs.rs index 46f0932ea..f9398e6d3 100644 --- a/dfdx/src/nn/layers/abs.rs +++ b/dfdx/src/nn/layers/abs.rs @@ -7,8 +7,7 @@ impl, T: Tape> crate::nn::Module; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_abs() } } diff --git a/dfdx/src/nn/layers/add_into.rs b/dfdx/src/nn/layers/add_into.rs index e1d46629f..983103396 100644 --- a/dfdx/src/nn/layers/add_into.rs +++ b/dfdx/src/nn/layers/add_into.rs @@ -31,7 +31,7 @@ pub struct AddInto( impl, T: BuildOnDevice> BuildOnDevice for AddInto { type Built = AddInto; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let t = self.0.try_build_on_device(device)?; Ok(AddInto(t)) } @@ -45,17 +45,16 @@ macro_rules! sum { macro_rules! add_into_impls { ($([$Mod:tt $ModVar:tt $Inp:tt $InpVar:tt]),+) => { impl< - Out: TryAdd, + Out: TryAdd, Ai, $($Inp, )+ A: Module, - $($Mod: Module<$Inp, Output = Out, Error = A::Error>, )+ + $($Mod: Module<$Inp, Output = Out>, )+ > Module<(Ai, $($Inp, )+)> for AddInto<(A, $($Mod, )+)> { type Output = Out; - type Error = A::Error; #[allow(clippy::needless_question_mark)] - fn try_forward(&self, x: (Ai, $($Inp, )+)) -> Result { + fn try_forward(&self, x: (Ai, $($Inp, )+)) -> Result { let (a, $($ModVar, )+) = &self.0; let (a_i, $($InpVar, )+) = x; let a_i = a.try_forward(a_i)?; @@ -63,7 +62,7 @@ macro_rules! add_into_impls { Ok(sum!(a_i, $($InpVar),*)) } #[allow(clippy::needless_question_mark)] - fn try_forward_mut(&mut self, x: (Ai, $($Inp, )+)) -> Result { + fn try_forward_mut(&mut self, x: (Ai, $($Inp, )+)) -> Result { let (a, $($ModVar, )+) = &mut self.0; let (a_i, $($InpVar, )+) = x; let a_i = a.try_forward_mut(a_i)?; diff --git a/dfdx/src/nn/layers/batch_norm1d.rs b/dfdx/src/nn/layers/batch_norm1d.rs index e0cf2da07..e6186365d 100644 --- a/dfdx/src/nn/layers/batch_norm1d.rs +++ b/dfdx/src/nn/layers/batch_norm1d.rs @@ -42,7 +42,7 @@ pub type BatchNorm1DConstConfig = BatchNorm1DConfig>; impl> BuildOnDevice for BatchNorm1DConfig { type Built = BatchNorm1D; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(BatchNorm1D { scale: device.try_ones_like(&(self.0,))?, bias: device.try_zeros_like(&(self.0,))?, @@ -82,7 +82,7 @@ pub struct BatchNorm1D> { } impl> crate::nn::ResetParams for BatchNorm1D { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { self.scale.try_fill_with_ones()?; self.bias.try_fill_with_zeros()?; self.running_mean.try_fill_with_zeros()?; @@ -94,12 +94,11 @@ impl, T: Tape> crate::nn::Module> for BatchNorm1D { type Output = Tensor<(B, C), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(B, C), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C), E, D, T>) -> Result { assert!(!T::OWNS_TAPE); self.infer_fwd(x) } - fn try_forward_mut(&mut self, x: Tensor<(B, C), E, D, T>) -> Result { + fn try_forward_mut(&mut self, x: Tensor<(B, C), E, D, T>) -> Result { assert!(T::OWNS_TAPE); self.train_fwd(x) } @@ -109,15 +108,11 @@ impl, T: Tape> crate::nn::Module> for BatchNorm1D { type Output = Tensor<(B, C, L), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(B, C, L), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C, L), E, D, T>) -> Result { assert!(!T::OWNS_TAPE); self.infer_fwd(x) } - fn try_forward_mut( - &mut self, - x: Tensor<(B, C, L), E, D, T>, - ) -> Result { + fn try_forward_mut(&mut self, x: Tensor<(B, C, L), E, D, T>) -> Result { assert!(T::OWNS_TAPE); self.train_fwd(x) } @@ -128,7 +123,7 @@ impl> BatchNorm1D { fn train_fwd, Ax: Axes>( &mut self, x: Tensor, - ) -> Result, D::Err> + ) -> Result, crate::tensor::Error> where S: HasAxes + ReduceShapeTo<(C,), Ax>, { @@ -175,7 +170,7 @@ impl> BatchNorm1D { pub fn infer_fwd, Ax: Axes>( &self, x: Tensor, - ) -> Result, D::Err> + ) -> Result, crate::tensor::Error> where (C,): BroadcastShapeTo, { diff --git a/dfdx/src/nn/layers/batch_norm2d.rs b/dfdx/src/nn/layers/batch_norm2d.rs index 40ca6243b..c6f592d38 100644 --- a/dfdx/src/nn/layers/batch_norm2d.rs +++ b/dfdx/src/nn/layers/batch_norm2d.rs @@ -44,7 +44,7 @@ pub type BatchNorm2DConstConfig = BatchNorm2DConfig>; impl> crate::nn::BuildOnDevice for BatchNorm2DConfig { type Built = BatchNorm2D; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(BatchNorm2D { scale: device.try_ones_like(&(self.0,))?, bias: device.try_zeros_like(&(self.0,))?, @@ -76,7 +76,7 @@ pub struct BatchNorm2D> { } impl> crate::nn::ResetParams for BatchNorm2D { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { self.scale.try_fill_with_ones()?; self.bias.try_fill_with_zeros()?; self.running_mean.try_fill_with_zeros()?; @@ -88,15 +88,11 @@ impl, T: Tape> crate::nn::Module> for BatchNorm2D { type Output = Tensor<(C, H, W), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(C, H, W), E, D, T>) -> Result { assert!(!T::OWNS_TAPE); self.infer_fwd(x) } - fn try_forward_mut( - &mut self, - x: Tensor<(C, H, W), E, D, T>, - ) -> Result { + fn try_forward_mut(&mut self, x: Tensor<(C, H, W), E, D, T>) -> Result { assert!(T::OWNS_TAPE); self.train_fwd(x) } @@ -106,18 +102,14 @@ impl, T: Tape> crate::nn::Module> for BatchNorm2D { type Output = Tensor<(Batch, C, H, W), E, D, T>; - type Error = D::Err; - fn try_forward( - &self, - x: Tensor<(Batch, C, H, W), E, D, T>, - ) -> Result { + fn try_forward(&self, x: Tensor<(Batch, C, H, W), E, D, T>) -> Result { assert!(!T::OWNS_TAPE); self.infer_fwd(x) } fn try_forward_mut( &mut self, x: Tensor<(Batch, C, H, W), E, D, T>, - ) -> Result { + ) -> Result { assert!(T::OWNS_TAPE); self.train_fwd(x) } @@ -128,7 +120,7 @@ impl> BatchNorm2D { fn train_fwd, Ax: Axes>( &mut self, x: Tensor, - ) -> Result, D::Err> + ) -> Result, crate::tensor::Error> where S: HasAxes + ReduceShapeTo<(C,), Ax>, { @@ -175,7 +167,7 @@ impl> BatchNorm2D { pub fn infer_fwd, Ax: Axes>( &self, x: Tensor, - ) -> Result, D::Err> + ) -> Result, crate::tensor::Error> where (C,): BroadcastShapeTo, { diff --git a/dfdx/src/nn/layers/bias1d.rs b/dfdx/src/nn/layers/bias1d.rs index feeed96fd..eb68a67bb 100644 --- a/dfdx/src/nn/layers/bias1d.rs +++ b/dfdx/src/nn/layers/bias1d.rs @@ -28,7 +28,7 @@ pub type Bias1DConstConfig = Bias1DConfig>; impl> BuildOnDevice for Bias1DConfig { type Built = Bias1D; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(Bias1D { bias: device.try_zeros_like(&(self.0,))?, }) @@ -44,7 +44,7 @@ pub struct Bias1D> { } impl> ResetParams for Bias1D { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { self.bias.try_fill_with_zeros() } } @@ -53,8 +53,7 @@ impl, T: Tape> Module for Bias1D { type Output = Tensor<(I,), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(I,), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(I,), E, D, T>) -> Result { x.try_add(self.bias.clone()) } } @@ -63,8 +62,7 @@ impl, T: Tape> Module { type Output = Tensor<(Batch, I), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(Batch, I), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(Batch, I), E, D, T>) -> Result { self.bias.retaped::().broadcast_like(&x).try_add(x) } } @@ -73,11 +71,7 @@ impl, T: Tape> Module> for Bias1D { type Output = Tensor<(Batch, Seq, I), E, D, T>; - type Error = D::Err; - fn try_forward( - &self, - x: Tensor<(Batch, Seq, I), E, D, T>, - ) -> Result { + fn try_forward(&self, x: Tensor<(Batch, Seq, I), E, D, T>) -> Result { self.bias.retaped::().broadcast_like(&x).try_add(x) } } diff --git a/dfdx/src/nn/layers/bias2d.rs b/dfdx/src/nn/layers/bias2d.rs index 497ceb196..1b6ac42ac 100644 --- a/dfdx/src/nn/layers/bias2d.rs +++ b/dfdx/src/nn/layers/bias2d.rs @@ -28,7 +28,7 @@ pub type Bias2DConstConfig = Bias2DConfig>; impl> BuildOnDevice for Bias2DConfig { type Built = Bias2D; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(Bias2D { bias: device.try_zeros_like(&(self.0,))?, }) @@ -44,7 +44,7 @@ pub struct Bias2D> { } impl> ResetParams for Bias2D { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { self.bias.try_fill_with_zeros() } } @@ -53,8 +53,7 @@ impl, T: Tape> Module> for Bias2D { type Output = Tensor<(C, H, W), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(C, H, W), E, D, T>) -> Result { self.bias.retaped::().broadcast_like(&x).try_add(x) } } @@ -63,8 +62,7 @@ impl, T: Tape> Module> for Bias2D { type Output = Tensor<(B, C, H, W), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(B, C, H, W), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C, H, W), E, D, T>) -> Result { self.bias.retaped::().broadcast_like(&x).try_add(x) } } diff --git a/dfdx/src/nn/layers/conv1d.rs b/dfdx/src/nn/layers/conv1d.rs index ba69e7672..793b326a9 100644 --- a/dfdx/src/nn/layers/conv1d.rs +++ b/dfdx/src/nn/layers/conv1d.rs @@ -62,7 +62,7 @@ where >::Output: Dim, { type Built = Conv1D; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { assert_eq!(self.in_chan.size() % self.groups.size(), 0); assert_eq!(self.out_chan.size() % self.groups.size(), 0); let i_over_g = self.in_chan / self.groups; @@ -119,7 +119,7 @@ where E: Dtype + num_traits::Float + rand_distr::uniform::SampleUniform, D: Device, { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { let (_, i_over_g, k) = self.weight.shape(); let scale = (1.0 / (k.size() * i_over_g.size()) as f64).sqrt(); let b = E::from_f64(scale).unwrap(); @@ -143,14 +143,7 @@ where L, G, >>::Convolved; - type Error = <(Img, Tensor<(O, >::Output, K), E, D>) as TryConv1D< - S, - P, - L, - G, - >>::Error; - - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { (x, self.weight.clone()).try_conv1d(self.stride, self.padding, self.dilation, self.groups) } } diff --git a/dfdx/src/nn/layers/conv2d.rs b/dfdx/src/nn/layers/conv2d.rs index 28a2552d6..9e65f41ff 100644 --- a/dfdx/src/nn/layers/conv2d.rs +++ b/dfdx/src/nn/layers/conv2d.rs @@ -78,7 +78,7 @@ where >::Output: Dim, { type Built = Conv2D; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { assert_eq!(self.in_chan.size() % self.groups.size(), 0); assert_eq!(self.out_chan.size() % self.groups.size(), 0); let i_over_g = self.in_chan / self.groups; @@ -141,7 +141,7 @@ where E: Dtype + num_traits::Float + rand_distr::uniform::SampleUniform, D: Device, { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { let (_, i_over_g, k, _) = self.weight.shape(); let scale = E::from_f64(1.0 / (k.size() * k.size() * i_over_g.size()) as f64).unwrap(); let b = scale.sqrt(); @@ -166,12 +166,7 @@ where Img, Tensor<(O, >::Output, K, K), E, D>, ) as TryConv2D>::Convolved; - type Error = <( - Img, - Tensor<(O, >::Output, K, K), E, D>, - ) as TryConv2D>::Error; - - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { (x, self.weight.clone()).try_conv2d(self.stride, self.padding, self.dilation, self.groups) } } diff --git a/dfdx/src/nn/layers/conv_trans2d.rs b/dfdx/src/nn/layers/conv_trans2d.rs index bfc2bf822..b4730a924 100644 --- a/dfdx/src/nn/layers/conv_trans2d.rs +++ b/dfdx/src/nn/layers/conv_trans2d.rs @@ -60,7 +60,7 @@ where >::Output: Dim, { type Built = ConvTrans2D; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { assert_eq!(self.in_chan.size() % self.groups.size(), 0); assert_eq!(self.out_chan.size() % self.groups.size(), 0); let o_over_g = self.out_chan / self.groups; @@ -119,7 +119,7 @@ where E: Dtype + num_traits::Float + rand_distr::uniform::SampleUniform, D: Device, { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { let (_, o_over_g, k, _) = self.weight.shape(); let b = (1.0 / (k.size() * k.size() * o_over_g.size()) as f64).sqrt(); let b = E::from_f64(b).unwrap(); @@ -144,12 +144,7 @@ where Img, Tensor<(I, >::Output, K, K), E, D>, ) as TryConvTrans2D>::Convolved; - type Error = <( - Img, - Tensor<(I, >::Output, K, K), E, D>, - ) as TryConvTrans2D>::Error; - - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { (x, self.weight.clone()).try_convtrans2d( self.stride, self.padding, diff --git a/dfdx/src/nn/layers/cos.rs b/dfdx/src/nn/layers/cos.rs index c0aa5fa5c..d9abf4e26 100644 --- a/dfdx/src/nn/layers/cos.rs +++ b/dfdx/src/nn/layers/cos.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Cos; impl, T: Tape> Module> for Cos { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_cos() } } diff --git a/dfdx/src/nn/layers/dropout.rs b/dfdx/src/nn/layers/dropout.rs index 741623ec3..a446ff912 100644 --- a/dfdx/src/nn/layers/dropout.rs +++ b/dfdx/src/nn/layers/dropout.rs @@ -23,10 +23,9 @@ impl, T: Tape> Module { type Output = Tensor; - type Error = D::Err; /// Does nothing - fn try_forward(&self, input: Tensor) -> Result { + fn try_forward(&self, input: Tensor) -> Result { assert!( !T::OWNS_TAPE, "DropoutOneIn::try_forward input must not be traced." @@ -35,7 +34,7 @@ impl, T: Tape> Module) -> Result { + fn try_forward_mut(&mut self, x: Tensor) -> Result { assert!( T::OWNS_TAPE, "DropoutOneIn::try_forward_mut input must be traced." @@ -70,10 +69,9 @@ impl Default for Dropout { impl, T: Tape> Module> for Dropout { type Output = Tensor; - type Error = D::Err; /// Does nothing - fn try_forward(&self, input: Tensor) -> Result { + fn try_forward(&self, input: Tensor) -> Result { assert!( !T::OWNS_TAPE, "Dropout::try_forward input must not be traced." @@ -82,7 +80,7 @@ impl, T: Tape> Module> } /// Applies dropout to the input tensor. - fn try_forward_mut(&mut self, x: Tensor) -> Result { + fn try_forward_mut(&mut self, x: Tensor) -> Result { assert!( T::OWNS_TAPE, "Dropout::try_forward_mut input must be traced." diff --git a/dfdx/src/nn/layers/embedding.rs b/dfdx/src/nn/layers/embedding.rs index 9f7767a48..2979c9d70 100644 --- a/dfdx/src/nn/layers/embedding.rs +++ b/dfdx/src/nn/layers/embedding.rs @@ -43,7 +43,7 @@ pub type EmbeddingConstConfig = impl> BuildOnDevice for EmbeddingConfig { type Built = Embedding; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(Embedding { weight: device.try_zeros_like(&(self.vocab, self.model))?, }) @@ -62,7 +62,7 @@ impl> ResetParams for Embedding, { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { self.weight.try_fill_with_distr(rand_distr::StandardNormal) } } @@ -71,9 +71,11 @@ impl, T: Tape> Module> for Embedding { type Output = Tensor<(Seq, M), E, D, T>; - type Error = D::Err; - fn try_forward(&self, input: Tensor<(Seq,), usize, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(Seq,), usize, D, T>, + ) -> Result { let (input, tape) = input.split_tape(); self.weight.clone().put_tape(tape).try_gather(input) } @@ -83,12 +85,11 @@ impl, T: Tape Module> for Embedding { type Output = Tensor<(Batch, Seq, M), E, D, T>; - type Error = D::Err; fn try_forward( &self, input: Tensor<(Batch, Seq), usize, D, T>, - ) -> Result { + ) -> Result { let (input, tape) = input.split_tape(); self.weight.clone().put_tape(tape).try_gather(input) } diff --git a/dfdx/src/nn/layers/exp.rs b/dfdx/src/nn/layers/exp.rs index a34c9b6fa..fd8b02492 100644 --- a/dfdx/src/nn/layers/exp.rs +++ b/dfdx/src/nn/layers/exp.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Exp; impl, T: Tape> Module> for Exp { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_exp() } } diff --git a/dfdx/src/nn/layers/flatten2d.rs b/dfdx/src/nn/layers/flatten2d.rs index a47b18875..bc414e3d9 100644 --- a/dfdx/src/nn/layers/flatten2d.rs +++ b/dfdx/src/nn/layers/flatten2d.rs @@ -14,9 +14,11 @@ where <>::Output as Mul>::Output: Dim, { type Output = Tensor<(<>::Output as Mul>::Output,), E, D, T>; - type Error = D::Err; - fn try_forward(&self, input: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(C, H, W), E, D, T>, + ) -> Result { let (c, h, w) = *input.shape(); let dst = (c * h * w,); input.try_reshape_like(&dst) @@ -31,12 +33,11 @@ where <>::Output as Mul>::Output: Dim, { type Output = Tensor<(Batch, <>::Output as Mul>::Output), E, D, T>; - type Error = D::Err; fn try_forward( &self, input: Tensor<(Batch, C, H, W), E, D, T>, - ) -> Result { + ) -> Result { let (batch, c, h, w) = *input.shape(); let dst = (batch, c * h * w); input.try_reshape_like(&dst) diff --git a/dfdx/src/nn/layers/gelu.rs b/dfdx/src/nn/layers/gelu.rs index 673094ead..4fa2188d0 100644 --- a/dfdx/src/nn/layers/gelu.rs +++ b/dfdx/src/nn/layers/gelu.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct FastGeLU; impl, T: Tape> Module> for FastGeLU { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_fast_gelu() } } @@ -16,8 +15,7 @@ impl, T: Tape> Module> pub struct AccurateGeLU; impl, T: Tape> Module> for AccurateGeLU { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_accurate_gelu() } } diff --git a/dfdx/src/nn/layers/generalized_add.rs b/dfdx/src/nn/layers/generalized_add.rs index 6ab0b8886..fb60b9c37 100644 --- a/dfdx/src/nn/layers/generalized_add.rs +++ b/dfdx/src/nn/layers/generalized_add.rs @@ -34,27 +34,25 @@ impl, T: BuildOnDevice, U: BuildOnDevice> Bui for GeneralizedAdd { type Built = GeneralizedAdd; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let t = self.t.try_build_on_device(device)?; let u = self.u.try_build_on_device(device)?; Ok(GeneralizedAdd { t, u }) } } -impl, U: Module> Module - for GeneralizedAdd +impl, U: Module> Module for GeneralizedAdd where - T::Output: TryAdd, + T::Output: TryAdd, { type Output = >::Output; - type Error = T::Error; - fn try_forward(&self, x: X) -> Result { + fn try_forward(&self, x: X) -> Result { let t = self.t.try_forward(x.with_empty_tape())?; let u = self.u.try_forward(x)?; t.try_add(u) } - fn try_forward_mut(&mut self, x: X) -> Result { + fn try_forward_mut(&mut self, x: X) -> Result { let t = self.t.try_forward_mut(x.with_empty_tape())?; let u = self.u.try_forward_mut(x)?; t.try_add(u) diff --git a/dfdx/src/nn/layers/generalized_mul.rs b/dfdx/src/nn/layers/generalized_mul.rs index a253d5f08..b6b5b5bae 100644 --- a/dfdx/src/nn/layers/generalized_mul.rs +++ b/dfdx/src/nn/layers/generalized_mul.rs @@ -33,27 +33,25 @@ impl, T: BuildOnDevice, U: BuildOnDevice> Bui for GeneralizedMul { type Built = GeneralizedMul; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let t = self.t.try_build_on_device(device)?; let u = self.u.try_build_on_device(device)?; Ok(GeneralizedMul { t, u }) } } -impl, U: Module> Module - for GeneralizedMul +impl, U: Module> Module for GeneralizedMul where - T::Output: TryMul, + T::Output: TryMul, { type Output = >::Output; - type Error = T::Error; - fn try_forward(&self, x: X) -> Result { + fn try_forward(&self, x: X) -> Result { let t = self.t.try_forward(x.with_empty_tape())?; let u = self.u.try_forward(x)?; t.try_mul(u) } - fn try_forward_mut(&mut self, x: X) -> Result { + fn try_forward_mut(&mut self, x: X) -> Result { let t = self.t.try_forward_mut(x.with_empty_tape())?; let u = self.u.try_forward_mut(x)?; t.try_mul(u) diff --git a/dfdx/src/nn/layers/layer_norm1d.rs b/dfdx/src/nn/layers/layer_norm1d.rs index 1bdb1a381..a1381534a 100644 --- a/dfdx/src/nn/layers/layer_norm1d.rs +++ b/dfdx/src/nn/layers/layer_norm1d.rs @@ -28,7 +28,7 @@ pub type LayerNorm1DConstConfig = LayerNorm1DConfig>; impl> BuildOnDevice for LayerNorm1DConfig { type Built = LayerNorm1D; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(LayerNorm1D { gamma: device.try_ones_like(&(self.0,))?, beta: device.try_zeros_like(&(self.0,))?, @@ -51,7 +51,7 @@ pub struct LayerNorm1D> { } impl> ResetParams for LayerNorm1D { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { self.gamma.try_fill_with_ones()?; self.beta.try_fill_with_zeros() } @@ -61,8 +61,7 @@ impl, T: Tape> Module for LayerNorm1D { type Output = Tensor<(M,), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(M,), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(M,), E, D, T>) -> Result { x.try_normalize(self.epsilon)? .try_mul(self.gamma.clone())? .try_add(self.beta.clone()) @@ -73,8 +72,7 @@ impl, T: Tape> Module { type Output = Tensor<(Batch, M), E, D, T>; - type Error = D::Err; - fn try_forward(&self, x: Tensor<(Batch, M), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(Batch, M), E, D, T>) -> Result { let x = x.try_normalize::>(self.epsilon)?; let x = self.gamma.retaped::().broadcast_like(&x).try_mul(x)?; self.beta.retaped::().broadcast_like(&x).try_add(x) @@ -85,11 +83,7 @@ impl, T: Tape> Module> for LayerNorm1D { type Output = Tensor<(Batch, Seq, M), E, D, T>; - type Error = D::Err; - fn try_forward( - &self, - x: Tensor<(Batch, Seq, M), E, D, T>, - ) -> Result { + fn try_forward(&self, x: Tensor<(Batch, Seq, M), E, D, T>) -> Result { let x = x.try_normalize::>(self.epsilon)?; let x = self.gamma.retaped::().broadcast_like(&x).try_mul(x)?; self.beta.retaped::().broadcast_like(&x).try_add(x) diff --git a/dfdx/src/nn/layers/leaky_relu.rs b/dfdx/src/nn/layers/leaky_relu.rs index 27ee7c214..a4ab83aaf 100644 --- a/dfdx/src/nn/layers/leaky_relu.rs +++ b/dfdx/src/nn/layers/leaky_relu.rs @@ -12,8 +12,7 @@ impl Default for LeakyReLU { impl, T: Tape> Module> for LeakyReLU { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_prelu(E::from_f64(self.0).unwrap()) } } diff --git a/dfdx/src/nn/layers/linear.rs b/dfdx/src/nn/layers/linear.rs index 74d6c711c..ce63b63af 100644 --- a/dfdx/src/nn/layers/linear.rs +++ b/dfdx/src/nn/layers/linear.rs @@ -38,7 +38,7 @@ pub type LinearConstConfig = LinearConfig> BuildOnDevice for LinearConfig { type Built = Linear; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(Linear { weight: device.try_zeros_like(&(self.out, self.inp))?, bias: device.try_zeros_like(&(self.out,))?, @@ -61,7 +61,7 @@ impl> ResetParams for Linear where E: Dtype + num_traits::Float + rand_distr::uniform::SampleUniform, { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { let (_o, i) = self.weight.shape(); let b = E::from_f64(1.0 / (i.size() as f64).sqrt()).unwrap(); self.weight.try_fill_with_distr(Uniform::new(-b, b))?; @@ -72,15 +72,13 @@ where impl, T: Tape> Module> for Linear where - Tensor: TryMatMul, Err = D::Err>, - Bias1D: - Module< as TryMatMul>>::Output, Error = D::Err>, + Tensor: TryMatMul>, + Bias1D: Module< as TryMatMul>>::Output>, { type Output = as Module< as TryMatMul>>::Output, >>::Output; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { let weight = self.weight.retaped::().try_permute()?; let bias = Bias1D { bias: self.bias.clone(), diff --git a/dfdx/src/nn/layers/ln.rs b/dfdx/src/nn/layers/ln.rs index 8d7124845..a493a9fc7 100644 --- a/dfdx/src/nn/layers/ln.rs +++ b/dfdx/src/nn/layers/ln.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Ln; impl, T: Tape> Module> for Ln { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_ln() } } diff --git a/dfdx/src/nn/layers/log_softmax.rs b/dfdx/src/nn/layers/log_softmax.rs index 47c58e783..c99a20fb5 100644 --- a/dfdx/src/nn/layers/log_softmax.rs +++ b/dfdx/src/nn/layers/log_softmax.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct LogSoftmax; impl, T: Tape> Module> for LogSoftmax { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_log_softmax() } } diff --git a/dfdx/src/nn/layers/matmul.rs b/dfdx/src/nn/layers/matmul.rs index 7b1de6694..86a3fa52f 100644 --- a/dfdx/src/nn/layers/matmul.rs +++ b/dfdx/src/nn/layers/matmul.rs @@ -28,7 +28,7 @@ pub type MatMulConstConfig = MatMulConfig> BuildOnDevice for MatMulConfig { type Built = MatMul; - fn try_build_on_device(&self, device: &D) -> Result { + fn try_build_on_device(&self, device: &D) -> Result { Ok(MatMul { weight: device.try_zeros_like(&(self.out, self.inp))?, }) @@ -47,7 +47,7 @@ impl> ResetParams for MatMul where E: Dtype + num_traits::Float + rand_distr::uniform::SampleUniform, { - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), Error> { let (_o, i) = self.weight.shape(); let scale = E::from_f64(1.0 / (i.size() as f64).sqrt()).unwrap(); self.weight.try_fill_with_distr(Uniform::new(-scale, scale)) @@ -57,11 +57,10 @@ where impl, T: Tape> Module> for MatMul where - Tensor: TryMatMul, Err = D::Err>, + Tensor: TryMatMul>, { type Output = as TryMatMul>>::Output; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_matmul(self.weight.retaped::().try_permute()?) } } diff --git a/dfdx/src/nn/layers/multi_head_attention.rs b/dfdx/src/nn/layers/multi_head_attention.rs index 0ccde7b9f..1232b4337 100644 --- a/dfdx/src/nn/layers/multi_head_attention.rs +++ b/dfdx/src/nn/layers/multi_head_attention.rs @@ -64,7 +64,6 @@ where T: Tape, { type Output = Tensor<(S1, M), E, D, T>; - type Error = D::Err; /// Encoder-Decoder style self attention where one set of tensors is used for values and keys, and another is used for queries fn try_forward( @@ -74,7 +73,7 @@ where Tensor<(S2, M), E, D>, Tensor<(S2, M), E, D>, ), - ) -> Result { + ) -> Result { assert_eq!(k.shape().0, v.shape().0); let (s1, m) = *q.shape(); let s2 = k.shape().0; @@ -101,7 +100,6 @@ where T: Tape, { type Output = Tensor<(B, S1, M), E, D, T>; - type Error = D::Err; /// Batched Encoder-Decoder style self attention where one set of tensors is used for values and keys, and another is used for queries fn try_forward( @@ -111,7 +109,7 @@ where Tensor<(B, S2, M), E, D>, Tensor<(B, S2, M), E, D>, ), - ) -> Result { + ) -> Result { assert_eq!(q.shape().0, k.shape().0); assert_eq!(q.shape().0, v.shape().0); assert_eq!(k.shape().1, v.shape().1); @@ -153,12 +151,11 @@ where E: Dtype, D: Device, Src: SplitTape, - Self: Module<(Src, Src::NoTape, Src::NoTape), Output = Src, Error = D::Err>, + Self: Module<(Src, Src::NoTape, Src::NoTape), Output = Src>, { type Output = Src; - type Error = D::Err; - fn try_forward(&self, src: Src) -> Result { + fn try_forward(&self, src: Src) -> Result { let (src, tape) = src.split_tape(); self.try_forward((src.clone().put_tape(tape), src.clone(), src)) } diff --git a/dfdx/src/nn/layers/pool_2d_avg.rs b/dfdx/src/nn/layers/pool_2d_avg.rs index 39e2d33da..6a32e86df 100644 --- a/dfdx/src/nn/layers/pool_2d_avg.rs +++ b/dfdx/src/nn/layers/pool_2d_avg.rs @@ -32,9 +32,8 @@ impl> Module for AvgPool2D { type Output = Img::Pooled; - type Error = Img::Error; - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { x.try_pool2d( crate::tensor_ops::Pool2DKind::Avg, self.kernel_size, diff --git a/dfdx/src/nn/layers/pool_2d_max.rs b/dfdx/src/nn/layers/pool_2d_max.rs index 09fd1ca4e..8a9ac78fd 100644 --- a/dfdx/src/nn/layers/pool_2d_max.rs +++ b/dfdx/src/nn/layers/pool_2d_max.rs @@ -32,9 +32,7 @@ impl> Module for MaxPool2D { type Output = Img::Pooled; - type Error = Img::Error; - - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { x.try_pool2d( crate::tensor_ops::Pool2DKind::Max, self.kernel_size, diff --git a/dfdx/src/nn/layers/pool_2d_min.rs b/dfdx/src/nn/layers/pool_2d_min.rs index af0726a80..18a194e63 100644 --- a/dfdx/src/nn/layers/pool_2d_min.rs +++ b/dfdx/src/nn/layers/pool_2d_min.rs @@ -32,9 +32,7 @@ impl> Module for MinPool2D { type Output = Img::Pooled; - type Error = Img::Error; - - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { x.try_pool2d( crate::tensor_ops::Pool2DKind::Min, self.kernel_size, diff --git a/dfdx/src/nn/layers/pool_global_avg.rs b/dfdx/src/nn/layers/pool_global_avg.rs index 599e55f91..cad457535 100644 --- a/dfdx/src/nn/layers/pool_global_avg.rs +++ b/dfdx/src/nn/layers/pool_global_avg.rs @@ -23,9 +23,10 @@ impl, T: Tape> Module> for AvgPoolGlobal { type Output = Tensor<(C,), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, input: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(C, H, W), E, D, T>, + ) -> Result { input.try_mean() } } @@ -34,9 +35,10 @@ impl, T: Tape> Module> for AvgPoolGlobal { type Output = Tensor<(B, C), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, input: Tensor<(B, C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(B, C, H, W), E, D, T>, + ) -> Result { input.try_mean() } } diff --git a/dfdx/src/nn/layers/pool_global_max.rs b/dfdx/src/nn/layers/pool_global_max.rs index e0eea18b6..1bf6de3d4 100644 --- a/dfdx/src/nn/layers/pool_global_max.rs +++ b/dfdx/src/nn/layers/pool_global_max.rs @@ -23,9 +23,10 @@ impl, T: Tape> Module> for MaxPoolGlobal { type Output = Tensor<(C,), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, input: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(C, H, W), E, D, T>, + ) -> Result { input.try_max() } } @@ -34,9 +35,10 @@ impl, T: Tape> Module> for MaxPoolGlobal { type Output = Tensor<(B, C), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, input: Tensor<(B, C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(B, C, H, W), E, D, T>, + ) -> Result { input.try_max() } } diff --git a/dfdx/src/nn/layers/pool_global_min.rs b/dfdx/src/nn/layers/pool_global_min.rs index 71f80ced8..e002dd554 100644 --- a/dfdx/src/nn/layers/pool_global_min.rs +++ b/dfdx/src/nn/layers/pool_global_min.rs @@ -23,9 +23,10 @@ impl, T: Tape> Module> for MinPoolGlobal { type Output = Tensor<(C,), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, input: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(C, H, W), E, D, T>, + ) -> Result { input.try_min() } } @@ -34,9 +35,10 @@ impl, T: Tape> Module> for MinPoolGlobal { type Output = Tensor<(B, C), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, input: Tensor<(B, C, H, W), E, D, T>) -> Result { + fn try_forward( + &self, + input: Tensor<(B, C, H, W), E, D, T>, + ) -> Result { input.try_min() } } diff --git a/dfdx/src/nn/layers/prelu.rs b/dfdx/src/nn/layers/prelu.rs index ca5eab481..6af121519 100644 --- a/dfdx/src/nn/layers/prelu.rs +++ b/dfdx/src/nn/layers/prelu.rs @@ -12,7 +12,7 @@ impl Default for PReLUConfig { impl> BuildOnDevice for PReLUConfig { type Built = PReLU; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let a = device.try_tensor(E::from_f64(self.0).unwrap())?; Ok(PReLU { a }) } @@ -28,16 +28,14 @@ pub struct PReLU> { impl> ResetParams for PReLU { /// Does nothing. - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { Ok(()) } } impl, T: Tape> Module> for PReLU { type Output = Tensor; - type Error = D::Err; - - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { let a = self.a.retaped::().broadcast_like(&x); x.try_prelu(a) } diff --git a/dfdx/src/nn/layers/prelu1d.rs b/dfdx/src/nn/layers/prelu1d.rs index e4c362523..b1a405c11 100644 --- a/dfdx/src/nn/layers/prelu1d.rs +++ b/dfdx/src/nn/layers/prelu1d.rs @@ -18,7 +18,7 @@ impl Default for PReLU1DConfig { impl> BuildOnDevice for PReLU1DConfig { type Built = PReLU1D; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let a = device.try_ones_like(&(self.c,))?.try_mul(self.a)?; Ok(PReLU1D { a }) } @@ -34,7 +34,7 @@ pub struct PReLU1D> { impl> ResetParams for PReLU1D { /// Does nothing. - fn try_reset_params(&mut self) -> Result<(), D::Err> { + fn try_reset_params(&mut self) -> Result<(), crate::tensor::Error> { Ok(()) } } @@ -43,9 +43,7 @@ impl, T: Tape> Module for PReLU1D { type Output = Tensor<(C,), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, x: Tensor<(C,), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(C,), E, D, T>) -> Result { x.try_prelu(self.a.clone()) } } @@ -54,9 +52,7 @@ impl, T: Tape> Module { type Output = Tensor<(B, C), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, x: Tensor<(B, C), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C), E, D, T>) -> Result { let a = self.a.retaped::().broadcast_like(&x); x.try_prelu(a) } @@ -66,9 +62,7 @@ impl, T: Tape> Module> for PReLU1D { type Output = Tensor<(B, C, H), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, x: Tensor<(B, C, H), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C, H), E, D, T>) -> Result { let a = self.a.retaped::().broadcast_like(&x); x.try_prelu(a) } @@ -78,9 +72,7 @@ impl, T: Tape> Module> for PReLU1D { type Output = Tensor<(B, C, H, W), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, x: Tensor<(B, C, H, W), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C, H, W), E, D, T>) -> Result { let a = self.a.retaped::().broadcast_like(&x); x.try_prelu(a) } diff --git a/dfdx/src/nn/layers/relu.rs b/dfdx/src/nn/layers/relu.rs index fee9cded8..d753ff7a9 100644 --- a/dfdx/src/nn/layers/relu.rs +++ b/dfdx/src/nn/layers/relu.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct ReLU; impl, T: Tape> Module> for ReLU { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_relu() } } diff --git a/dfdx/src/nn/layers/reshape.rs b/dfdx/src/nn/layers/reshape.rs index 83be2344a..83d977bab 100644 --- a/dfdx/src/nn/layers/reshape.rs +++ b/dfdx/src/nn/layers/reshape.rs @@ -18,8 +18,7 @@ impl, T: Tape> Module { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_reshape_like(&self.0) } } diff --git a/dfdx/src/nn/layers/residual_add.rs b/dfdx/src/nn/layers/residual_add.rs index d9c707167..5da249a9f 100644 --- a/dfdx/src/nn/layers/residual_add.rs +++ b/dfdx/src/nn/layers/residual_add.rs @@ -27,10 +27,9 @@ pub struct ResidualAdd( pub T, ); -// TODO derive this impl, T: BuildOnDevice> BuildOnDevice for ResidualAdd { type Built = ResidualAdd; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let t = self.0.try_build_on_device(device)?; Ok(ResidualAdd(t)) } @@ -38,15 +37,14 @@ impl, T: BuildOnDevice> BuildOnDevice for Res impl> Module for ResidualAdd where - T::Output: TryAdd, + T::Output: TryAdd, { type Output = >::Output; - type Error = T::Error; - fn try_forward(&self, x: X) -> Result { + fn try_forward(&self, x: X) -> Result { let t = self.0.try_forward(x.with_empty_tape())?; t.try_add(x) } - fn try_forward_mut(&mut self, x: X) -> Result { + fn try_forward_mut(&mut self, x: X) -> Result { let t = self.0.try_forward_mut(x.with_empty_tape())?; t.try_add(x) } diff --git a/dfdx/src/nn/layers/residual_mul.rs b/dfdx/src/nn/layers/residual_mul.rs index 428c1c214..c55787ffa 100644 --- a/dfdx/src/nn/layers/residual_mul.rs +++ b/dfdx/src/nn/layers/residual_mul.rs @@ -26,10 +26,9 @@ pub struct ResidualMul( pub T, ); -// TODO derive this impl, T: BuildOnDevice> BuildOnDevice for ResidualMul { type Built = ResidualMul; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let t = self.0.try_build_on_device(device)?; Ok(ResidualMul(t)) } @@ -37,15 +36,14 @@ impl, T: BuildOnDevice> BuildOnDevice for Res impl> Module for ResidualMul where - T::Output: TryMul, + T::Output: TryMul, { type Output = >::Output; - type Error = T::Error; - fn try_forward(&self, x: X) -> Result { + fn try_forward(&self, x: X) -> Result { let t = self.0.try_forward(x.with_empty_tape())?; t.try_mul(x) } - fn try_forward_mut(&mut self, x: X) -> Result { + fn try_forward_mut(&mut self, x: X) -> Result { let t = self.0.try_forward_mut(x.with_empty_tape())?; t.try_mul(x) } diff --git a/dfdx/src/nn/layers/sigmoid.rs b/dfdx/src/nn/layers/sigmoid.rs index 56b09d74a..3eef17462 100644 --- a/dfdx/src/nn/layers/sigmoid.rs +++ b/dfdx/src/nn/layers/sigmoid.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Sigmoid; impl, T: Tape> Module> for Sigmoid { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_sigmoid() } } diff --git a/dfdx/src/nn/layers/sin.rs b/dfdx/src/nn/layers/sin.rs index 57074fd1c..eae38f486 100644 --- a/dfdx/src/nn/layers/sin.rs +++ b/dfdx/src/nn/layers/sin.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Sin; impl, T: Tape> Module> for Sin { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_sin() } } diff --git a/dfdx/src/nn/layers/softmax.rs b/dfdx/src/nn/layers/softmax.rs index 18a28544a..145a054b2 100644 --- a/dfdx/src/nn/layers/softmax.rs +++ b/dfdx/src/nn/layers/softmax.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Softmax; impl, T: Tape> Module> for Softmax { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_softmax() } } diff --git a/dfdx/src/nn/layers/split_into.rs b/dfdx/src/nn/layers/split_into.rs index 1e66d68fe..de905a0df 100644 --- a/dfdx/src/nn/layers/split_into.rs +++ b/dfdx/src/nn/layers/split_into.rs @@ -33,7 +33,7 @@ pub struct SplitInto( impl, T: BuildOnDevice> BuildOnDevice for SplitInto { type Built = SplitInto; - fn try_build_on_device(&self, device: &D) -> Result::Err> { + fn try_build_on_device(&self, device: &D) -> Result { let t = self.0.try_build_on_device(device)?; Ok(SplitInto(t)) } @@ -44,29 +44,26 @@ macro_rules! tuple_impls { impl< Input: WithEmptyTape, $head: Module, - $($tails : Module,)+ + $($tails : Module,)+ > Module for SplitInto<($head, $($tails,)+)> { type Output = ( $head::Output, $($tails::Output),+ ); - type Error = $head::Error; #[allow(non_snake_case)] - fn try_forward(&self, x: Input) -> Result { + fn try_forward(&self, x: Input) -> Result { let ($head, $($tails,)+) = &self.0; let ($($tails,)+) = ($($tails.try_forward(x.with_empty_tape())?,)+); let $head = $head.try_forward(x)?; - Ok(($head, $($tails,)+)) } #[allow(non_snake_case)] - fn try_forward_mut(&mut self, x: Input) -> Result { + fn try_forward_mut(&mut self, x: Input) -> Result { let ($head, $($tails,)+) = &mut self.0; let ($($tails,)+) = ($($tails.try_forward_mut(x.with_empty_tape())?,)+); let $head = $head.try_forward_mut(x)?; - Ok(($head, $($tails,)+)) } } diff --git a/dfdx/src/nn/layers/sqrt.rs b/dfdx/src/nn/layers/sqrt.rs index cb9246bf8..4621b7092 100644 --- a/dfdx/src/nn/layers/sqrt.rs +++ b/dfdx/src/nn/layers/sqrt.rs @@ -6,8 +6,7 @@ pub struct Sqrt; impl, T: Tape> Module> for Sqrt { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_sqrt() } } diff --git a/dfdx/src/nn/layers/square.rs b/dfdx/src/nn/layers/square.rs index b9e030d39..5b2ac5d51 100644 --- a/dfdx/src/nn/layers/square.rs +++ b/dfdx/src/nn/layers/square.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Square; impl, T: Tape> Module> for Square { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_square() } } diff --git a/dfdx/src/nn/layers/tanh.rs b/dfdx/src/nn/layers/tanh.rs index bba87199e..6507f5807 100644 --- a/dfdx/src/nn/layers/tanh.rs +++ b/dfdx/src/nn/layers/tanh.rs @@ -5,8 +5,7 @@ use crate::prelude::*; pub struct Tanh; impl, T: Tape> Module> for Tanh { type Output = Tensor; - type Error = D::Err; - fn try_forward(&self, x: Tensor) -> Result { + fn try_forward(&self, x: Tensor) -> Result { x.try_tanh() } } diff --git a/dfdx/src/nn/layers/transformer.rs b/dfdx/src/nn/layers/transformer.rs index e4a985a12..fa7ab76a3 100644 --- a/dfdx/src/nn/layers/transformer.rs +++ b/dfdx/src/nn/layers/transformer.rs @@ -100,17 +100,15 @@ impl DecoderBlockConfig { impl, Tgt, Mem> Module<(Tgt, Mem)> for DecoderBlock where - Tgt: WithEmptyTape + SplitTape + TryAdd + HasErr, + Tgt: WithEmptyTape + SplitTape + TryAdd, Mem: Clone, - ResidualAdd>: Module, - MultiHeadAttention: Module<(Tgt, Mem, Mem), Output = Tgt, Error = D::Err>, - LayerNorm1D: Module, - ResidualAdd>: Module, + ResidualAdd>: Module, + MultiHeadAttention: Module<(Tgt, Mem, Mem), Output = Tgt>, + LayerNorm1D: Module, + ResidualAdd>: Module, { type Output = Tgt; - type Error = D::Err; - - fn try_forward(&self, (tgt, mem): (Tgt, Mem)) -> Result { + fn try_forward(&self, (tgt, mem): (Tgt, Mem)) -> Result { let x = self.self_attn.try_forward(tgt)?; let x = self.norm1.try_forward(x)?; @@ -179,17 +177,14 @@ impl TransformerConfig { impl, Src: SplitTape, Tgt: PutTape> Module<(Src, Tgt)> for Transformer where - Vec>: Module, + Vec>: Module, DecoderBlock: Module< (>::Output, Src::NoTape), Output = >::Output, - Error = D::Err, >, { type Output = >::Output; - type Error = D::Err; - - fn try_forward(&self, (src, tgt): (Src, Tgt)) -> Result { + fn try_forward(&self, (src, tgt): (Src, Tgt)) -> Result { let (mem, tape) = self.encoder.try_forward(src)?.split_tape(); let mut tgt = tgt.put_tape(tape); for block in self.decoder.iter() { diff --git a/dfdx/src/nn/layers/upscale2d.rs b/dfdx/src/nn/layers/upscale2d.rs index d8bdd31d7..33f28cd64 100644 --- a/dfdx/src/nn/layers/upscale2d.rs +++ b/dfdx/src/nn/layers/upscale2d.rs @@ -18,9 +18,7 @@ impl> Module for Upscale2D { type Output = Img::Output; - type Error = Img::Err; - - fn try_forward(&self, x: Img) -> Result { + fn try_forward(&self, x: Img) -> Result { x.generic_upscale2d_like(self.method, self.out_height, self.out_width) } } @@ -45,9 +43,7 @@ where D: Device + Upscale2DKernel, { type Output = Tensor<(C, H::Output, W::Output), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, x: Tensor<(C, H, W), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(C, H, W), E, D, T>) -> Result { let (_c, h, w) = *x.shape(); let h = h * self.height_factor; let w = w * self.width_factor; @@ -66,9 +62,7 @@ where T: 'static + Tape, { type Output = Tensor<(B, C, H::Output, W::Output), E, D, T>; - type Error = D::Err; - - fn try_forward(&self, x: Tensor<(B, C, H, W), E, D, T>) -> Result { + fn try_forward(&self, x: Tensor<(B, C, H, W), E, D, T>) -> Result { let (_b, _c, h, w) = *x.shape(); let h = h * self.height_factor; let w = w * self.width_factor; diff --git a/dfdx/src/nn/optim/adam.rs b/dfdx/src/nn/optim/adam.rs index 8fb4cc383..ff1d86b5f 100644 --- a/dfdx/src/nn/optim/adam.rs +++ b/dfdx/src/nn/optim/adam.rs @@ -2,7 +2,7 @@ use std::marker::PhantomData; use crate::{ shapes::{Dtype, Shape}, - tensor::{Gradients, Storage, Tensor, Tensorlike, UniqueId}, + tensor::{Error, Gradients, Storage, Tensor, Tensorlike, UniqueId}, tensor_ops::{AdamConfig, Device}, }; @@ -55,7 +55,7 @@ impl> crate::nn::Optimizer for Adam t: &mut Tensor, gradients: &Gradients, missing_params: &mut Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { let g = gradients.get_ref_checked(t); match g { None => missing_params.push(t.id()), @@ -68,11 +68,7 @@ impl> crate::nn::Optimizer for Adam Ok(()) } - fn update( - &mut self, - module: &mut M, - gradients: &Gradients, - ) -> Result<(), crate::nn::OptimizerUpdateError<::Err>> + fn update(&mut self, module: &mut M, gradients: &Gradients) -> Result<(), Error> where M: crate::nn::UpdateParams, { @@ -80,15 +76,11 @@ impl> crate::nn::Optimizer for Adam // NOTE: the rest of this is identical to default implementation of update. let mut missing_tensors = Vec::new(); - module - .try_update_params(self, gradients, &mut missing_tensors) - .map_err(crate::nn::OptimizerUpdateError::DeviceError)?; + module.try_update_params(self, gradients, &mut missing_tensors)?; if missing_tensors.is_empty() { Ok(()) } else { - Err(crate::nn::OptimizerUpdateError::UnusedTensors( - missing_tensors, - )) + Err(Error::UnusedTensors(missing_tensors)) } } } diff --git a/dfdx/src/nn/optim/mod.rs b/dfdx/src/nn/optim/mod.rs index 790b6fceb..e80e549ee 100644 --- a/dfdx/src/nn/optim/mod.rs +++ b/dfdx/src/nn/optim/mod.rs @@ -39,5 +39,5 @@ pub use adam::Adam; pub use rmsprop::RMSprop; pub use sgd::Sgd; // re-exports -pub use super::{Optimizer, OptimizerUpdateError}; +pub use super::Optimizer; pub use crate::tensor_ops::{AdamConfig, Momentum, RMSpropConfig, SgdConfig, WeightDecay}; diff --git a/dfdx/src/nn/optim/rmsprop.rs b/dfdx/src/nn/optim/rmsprop.rs index af7d6493b..92db198a3 100644 --- a/dfdx/src/nn/optim/rmsprop.rs +++ b/dfdx/src/nn/optim/rmsprop.rs @@ -2,7 +2,7 @@ use std::marker::PhantomData; use crate::{ shapes::{Dtype, Shape}, - tensor::{Gradients, Storage, Tensor, Tensorlike, UniqueId}, + tensor::{Error, Gradients, Storage, Tensor, Tensorlike, UniqueId}, tensor_ops::{Device, RMSpropConfig}, }; @@ -66,7 +66,7 @@ impl> crate::nn::Optimizer for RMSprop, gradients: &Gradients, missing_params: &mut Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), Error> { let g = gradients.get_ref_checked(t); match g { None => missing_params.push(t.id()), @@ -85,25 +85,17 @@ impl> crate::nn::Optimizer for RMSprop, - ) -> Result<(), crate::nn::OptimizerUpdateError<::Err>> + fn update(&mut self, module: &mut M, gradients: &Gradients) -> Result<(), Error> where M: crate::nn::UpdateParams, { // NOTE: the rest of this is identical to default implementation of update. let mut missing_tensors = Vec::new(); - module - .try_update_params(self, gradients, &mut missing_tensors) - .map_err(crate::nn::OptimizerUpdateError::DeviceError)?; + module.try_update_params(self, gradients, &mut missing_tensors)?; let r = if missing_tensors.is_empty() { Ok(()) } else { - Err(crate::nn::OptimizerUpdateError::UnusedTensors( - missing_tensors, - )) + Err(Error::UnusedTensors(missing_tensors)) }; self.step += 1; r diff --git a/dfdx/src/nn/optim/sgd.rs b/dfdx/src/nn/optim/sgd.rs index 33d00c7a0..9cad6e0e9 100644 --- a/dfdx/src/nn/optim/sgd.rs +++ b/dfdx/src/nn/optim/sgd.rs @@ -51,7 +51,7 @@ impl> crate::nn::Optimizer for Sgd { t: &mut Tensor, gradients: &Gradients, missing_params: &mut Vec, - ) -> Result<(), D::Err> { + ) -> Result<(), crate::tensor::Error> { let g = gradients.get_ref_checked(t); match g { None => missing_params.push(t.id()),