Merge branch 'main' into don/feat/wgpu_to_dtype

coreylowman · Jan 25, 2024 · 16658b2 · 16658b2
2 parents a1fcda8 + e04dd4f
commit 16658b2
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 5 deletions.
diff --git a/dfdx-core/src/tensor_ops/conv1d/mod.rs b/dfdx-core/src/tensor_ops/conv1d/mod.rs
@@ -9,7 +9,7 @@ mod tests;
 
 #[repr(C)]
 #[derive(Debug, Copy, Clone)]
-pub(super) struct Conv1DOp {
+pub struct Conv1DOp {
     pub kernel: usize,
     pub stride: usize,
     pub padding: usize,
@@ -22,7 +22,7 @@ pub(super) struct Conv1DOp {
     pub l_out: usize,
 }
 
-pub(super) trait Conv1DKernel<E: Dtype>: Storage<E> {
+pub trait Conv1DKernel<E: Dtype>: Storage<E> {
     fn alloc<S: Shape>(&self, s: S) -> Result<Tensor<S, E, Self>, Error>;
 
     fn forward<L: Shape, R: Shape, O: Shape>(
@@ -108,6 +108,7 @@ pub trait TryConv1D<Stride, Padding, Dilation, Groups>: Sized {
     ) -> Result<Self::Convolved, Error>;
 }
 
+#[cfg(feature = "nightly")]
 impl<
         const KERNEL: usize,
         const STRIDE: usize,

diff --git a/dfdx-core/src/tensor_ops/mod.rs b/dfdx-core/src/tensor_ops/mod.rs
@@ -281,9 +281,7 @@ pub use upscale2d::{
 };
 pub use var_to::VarTo;
 
-#[cfg(feature = "nightly")]
 mod conv1d;
-#[cfg(feature = "nightly")]
 pub use conv1d::TryConv1D;
 
 #[cfg(feature = "nightly")]

diff --git a/dfdx-core/src/tensor_ops/utilities/device.rs b/dfdx-core/src/tensor_ops/utilities/device.rs
@@ -112,6 +112,9 @@ pub trait Device<E: Dtype>:
     + BinaryKernel<super::super::maximum::MaximumKernelOp, E>
     + BinaryKernel<super::super::minimum::MinimumKernelOp, E>
     + crate::tensor_ops::axpy::AxpyKernel<E>
+
+    // conv1d
+    + super::super::conv1d::Conv1DKernel<E>
 {
 }
 

diff --git a/dfdx/src/lib.rs b/dfdx/src/lib.rs
@@ -224,7 +224,7 @@
 //! | Optimizer | dfdx | pytorch |
 //! | --- | --- | --- |
 //! | SGD | [nn::optim::Sgd] | `torch.optim.SGD` |
-//! | Adam | [nn::optim::Adam] | torch.optim.Adam` |
+//! | Adam | [nn::optim::Adam] | `torch.optim.Adam` |
 //! | AdamW | [nn::optim::Adam] with [nn::optim::WeightDecay::Decoupled] | `torch.optim.AdamW` |
 //! | RMSprop | [nn::optim::RMSprop] | `torch.optim.RMSprop` |
 //!