diff --git a/dfdx-core/src/data/arange.rs b/dfdx-core/src/data/arange.rs
index fbb44b9f2..ff106c58e 100644
--- a/dfdx-core/src/data/arange.rs
+++ b/dfdx-core/src/data/arange.rs
@@ -11,7 +11,7 @@ pub trait Arange<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E> {
     ///
     /// Const sized tensor:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::Arange};
+    /// # use dfdx_core::{prelude::*, data::Arange};
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank1<5>, f32, _> = dev.arange(Const::<5>);
     /// assert_eq!(t.array(), [0.0, 1.0, 2.0, 3.0, 4.0]);
@@ -19,7 +19,7 @@ pub trait Arange<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E> {
     ///
     /// Runtime sized tensor:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::Arange};
+    /// # use dfdx_core::{prelude::*, data::Arange};
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<(usize, ), f32, _> = dev.arange(5);
     /// assert_eq!(t.as_vec(), [0.0, 1.0, 2.0, 3.0, 4.0]);
diff --git a/dfdx-core/src/data/batch.rs b/dfdx-core/src/data/batch.rs
index 70e1f1285..7a3666c27 100644
--- a/dfdx-core/src/data/batch.rs
+++ b/dfdx-core/src/data/batch.rs
@@ -80,14 +80,14 @@ pub trait IteratorBatchExt: Iterator {
     ///
     /// Const batches:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::IteratorBatchExt};
+    /// # use dfdx_core::{prelude::*, data::IteratorBatchExt};
     /// let items: Vec<[usize; 5]> = (0..12).batch_exact(Const::<5>).collect();
     /// assert_eq!(&items, &[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]);
     /// ```
     ///
     /// Runtime batches:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::IteratorBatchExt};
+    /// # use dfdx_core::{prelude::*, data::IteratorBatchExt};
     /// let items: Vec<Vec<usize>> = (0..12).batch_exact(5).collect();
     /// assert_eq!(&items, &[[0, 1, 2, 3, 4], [5, 6, 7, 8, 9]]);
     /// ```
@@ -104,7 +104,7 @@ pub trait IteratorBatchExt: Iterator {
     ///
     /// Example:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::IteratorBatchExt};
+    /// # use dfdx_core::{prelude::*, data::IteratorBatchExt};
     /// let items: Vec<Vec<usize>> = (0..12).batch_with_last(5).collect();
     /// assert_eq!(&items, &[vec![0, 1, 2, 3, 4], vec![5, 6, 7, 8, 9], vec![10, 11]]);
     /// ```
diff --git a/dfdx-core/src/data/collate.rs b/dfdx-core/src/data/collate.rs
index c29706b06..d38a2a677 100644
--- a/dfdx-core/src/data/collate.rs
+++ b/dfdx-core/src/data/collate.rs
@@ -93,7 +93,7 @@ pub trait IteratorCollateExt: Iterator {
     ///
     /// Example implementations:
     /// ```rust
-    /// # use dfdx::data::IteratorCollateExt;
+    /// # use dfdx_core::data::IteratorCollateExt;
     /// let data = [[('a', 'b'); 10], [('c', 'd'); 10], [('e', 'f'); 10]];
     /// // we use collate to transform each batch:
     /// let mut iter = data.into_iter().collate();
diff --git a/dfdx-core/src/data/one_hot_encode.rs b/dfdx-core/src/data/one_hot_encode.rs
index 9347ad346..c3ce86107 100644
--- a/dfdx-core/src/data/one_hot_encode.rs
+++ b/dfdx-core/src/data/one_hot_encode.rs
@@ -16,7 +16,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Const class labels and const n:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = [0, 1, 2, 1, 1];
     /// let probs: Tensor<Rank2<5, 3>, f32, _> = dev.one_hot_encode(Const::<3>, class_labels);
@@ -31,7 +31,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Runtime class labels and const n:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = [0, 1, 2, 1, 1];
     /// let probs: Tensor<(Const<5>, usize), f32, _> = dev.one_hot_encode(3, class_labels);
@@ -46,7 +46,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Const class labels and runtime n:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = std::vec![0, 1, 2, 1, 1];
     /// let probs: Tensor<(usize, Const<3>), f32, _> = dev.one_hot_encode(Const, class_labels);
@@ -61,7 +61,7 @@ pub trait OneHotEncode<E: Dtype>: Storage<E> + ZerosTensor<E> + TensorFromVec<E>
     ///
     /// Runtime both:
     /// ```rust
-    /// # use dfdx::{prelude::*, data::OneHotEncode};
+    /// # use dfdx_core::{prelude::*, data::OneHotEncode};
     /// # let dev: Cpu = Default::default();
     /// let class_labels = std::vec![0, 1, 2, 1, 1];
     /// let probs: Tensor<(usize, usize), f32, _> = dev.one_hot_encode(3, class_labels);
diff --git a/dfdx-core/src/data/stack.rs b/dfdx-core/src/data/stack.rs
index 42f83b402..c751c3215 100644
--- a/dfdx-core/src/data/stack.rs
+++ b/dfdx-core/src/data/stack.rs
@@ -29,7 +29,7 @@ pub trait IteratorStackExt: Iterator {
     ///
     /// Example implementations:
     /// ```rust
-    /// # use dfdx::{data::IteratorStackExt, prelude::*};
+    /// # use dfdx_core::{data::IteratorStackExt, prelude::*};
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank1<3>, f32, _> = dev.zeros();
     /// let data = [[a.clone(), a.clone(), a]];
diff --git a/dfdx-core/src/lib.rs b/dfdx-core/src/lib.rs
index 337b43e8a..2dfab274b 100644
--- a/dfdx-core/src/lib.rs
+++ b/dfdx-core/src/lib.rs
@@ -67,7 +67,7 @@
 //! Here's how you might use a device:
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! let dev: Cpu = Default::default();
 //! let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
 //! ```
diff --git a/dfdx-core/src/shapes/mod.rs b/dfdx-core/src/shapes/mod.rs
index 03a2d91a4..36753780e 100644
--- a/dfdx-core/src/shapes/mod.rs
+++ b/dfdx-core/src/shapes/mod.rs
@@ -2,7 +2,7 @@
 //!
 //! Example shapes:
 //! ```rust
-//! # use dfdx::shapes::*;
+//! # use dfdx_core::shapes::*;
 //! let _: Rank3<2, 3, 4> = Default::default();
 //! let _: (Const<2>, Const<3>) = Default::default();
 //! let _: (usize, Const<4>) = (3, Const);
diff --git a/dfdx-core/src/tensor/mod.rs b/dfdx-core/src/tensor/mod.rs
index a6e79cbff..1b3a5b06d 100644
--- a/dfdx-core/src/tensor/mod.rs
+++ b/dfdx-core/src/tensor/mod.rs
@@ -15,13 +15,13 @@
 //! In order to do anything with tensors, you first need to construct the device that they will be stored on:
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! let dev: Cpu = Default::default();
 //! let dev: Cpu = Cpu::seed_from_u64(0);
 //! ```
 //!
 //! ```ignore
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! let dev: Cuda = Default::default();
 //! let dev: Cuda = Cuda::seed_from_u64(1234);
 //! let dev: Cuda = Cuda::try_build(0, 1234).unwrap();
@@ -34,7 +34,7 @@
 //! See [TensorFrom] & [TensorFromVec].
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let _ = dev.tensor([1.0, 2.0, 3.0]);
 //! let _ = dev.tensor_from_vec(vec![1.0, 2.0, 3.0], (3, ));
@@ -45,7 +45,7 @@
 //! See [ZerosTensor] and [OnesTensor].
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let _: Tensor<Rank1<5>,f32 , _> = dev.zeros();
 //! let _: Tensor<Rank2<3, 2>, f32, _> = dev.ones();
@@ -56,7 +56,7 @@
 //! See [SampleTensor]
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let _: Tensor<Rank1<5>, f32, _> = dev.sample_uniform();
 //! let _: Tensor<Rank2<3, 5>, f32, _> = dev.sample_normal();
@@ -70,7 +70,7 @@
 //! You can use [Tensor::copy_from] and [Tensor::copy_into] to copy data into a tensor:
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let mut a: Tensor<Rank1<1000>, f32, _> = dev.zeros();
 //! let buf: Vec<f32> = vec![1.0; 1000];
@@ -90,7 +90,7 @@
 //! with them directly.
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
 //! let t: [[f32; 3]; 2] = t.array();
@@ -105,7 +105,7 @@
 //! Note that these two methods are only present for tensors without a tape already.
 //!
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t: Tensor<Rank1<5>,f32, _> = dev.zeros();
 //! let mut grads = Gradients::leaky();
diff --git a/dfdx-core/src/tensor/storage_traits.rs b/dfdx-core/src/tensor/storage_traits.rs
index 4e3e33dfa..4f98c3e41 100644
--- a/dfdx-core/src/tensor/storage_traits.rs
+++ b/dfdx-core/src/tensor/storage_traits.rs
@@ -104,7 +104,7 @@ impl<S: Shape, E, D: CopySlice<E>, T> Tensor<S, E, D, T> {
     /// Copy *physical* data from a slice - **panics** if there are not enough elements in the slice.
     ///
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let data = [1.0, 2.0, 3.0, 4.0];
     /// let mut t: Tensor<Rank2<2, 2>, f32, _> = dev.zeros();
@@ -118,7 +118,7 @@ impl<S: Shape, E, D: CopySlice<E>, T> Tensor<S, E, D, T> {
     /// Copy *physical* data into a slice - **panics** if there are not enough elements in the tensor.
     ///
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank2<2, 2>, f32, _> = dev.tensor([[1.0, 2.0], [3.0, 4.0]]);
     /// let mut data = [0.0; 4];
@@ -134,7 +134,7 @@ impl<S: Shape, E, D: CopySlice<E>, T> Tensor<S, E, D, T> {
 pub trait ZerosTensor<E>: Storage<E> {
     /// Creates a tensor filled with zeros.
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
     /// ```
@@ -151,14 +151,14 @@ pub trait ZerosTensor<E>: Storage<E> {
     ///
     /// Given a shape directly:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(5, Const));
     /// ```
     ///
     /// Given another tensor:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
     /// let b: Tensor<Rank2<2, 3>, f32, _> = dev.zeros_like(&a);
@@ -179,7 +179,7 @@ pub trait ZeroFillStorage<E>: Storage<E> {
 pub trait OnesTensor<E>: Storage<E> {
     /// Creates a tensor filled with ones.
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.ones();
     /// ```
@@ -196,14 +196,14 @@ pub trait OnesTensor<E>: Storage<E> {
     ///
     /// Given a shape directly:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<(usize, Const<3>), f32, _> = dev.ones_like(&(5, Const));
     /// ```
     ///
     /// Given another tensor:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.ones();
     /// let b: Tensor<_, f32, _> = dev.ones_like(&a);
@@ -231,7 +231,7 @@ pub trait TriangleTensor<E>: Storage<E> {
     /// ## Examples
     ///
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<3, 3>, f32, _> = dev.upper_tri(1.0, None);
     /// assert_eq!(a.array(),
@@ -297,7 +297,7 @@ pub trait TriangleTensor<E>: Storage<E> {
     /// ## Examples
     ///
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<3, 3>, f32, _> = dev.lower_tri(1.0, None);
     /// assert_eq!(a.array(),
@@ -480,7 +480,7 @@ impl<S: Shape, E, D: Storage<E>, T> Tensor<S, E, D, T> {
 pub trait TensorFrom<Src, S: Shape, E>: Storage<E> {
     /// Create a tensor from rust data
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let _: Tensor<Rank2<2, 3>, f32, Cpu> = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
     /// let _: Tensor<Rank2<2, 3>, f32, Cpu> = dev.tensor(vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]);
diff --git a/dfdx-core/src/tensor/tensor_impls.rs b/dfdx-core/src/tensor/tensor_impls.rs
index 5d7a96674..bf4c57f22 100644
--- a/dfdx-core/src/tensor/tensor_impls.rs
+++ b/dfdx-core/src/tensor/tensor_impls.rs
@@ -17,7 +17,7 @@ use std::sync::Arc;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// // A 1d tensor with 1000 f32 elements, stored on the Cpu
 /// type A = Tensor<Rank1<1000>, f32, Cpu>;
@@ -116,7 +116,7 @@ impl<S: Shape, E, D: Storage<E>, T> Tensor<S, E, D, T> {
 pub trait PutTape<T> {
     type Output;
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let a: Tensor<Rank2<2, 3>, f32, _, NoneTape> = dev.zeros();
     /// let a: Tensor<Rank2<2, 3>, f32, _, OwnedTape<f32, Cpu>> = a.put_tape(Default::default());
@@ -146,7 +146,7 @@ pub trait SplitTape {
     type NoTape: Clone + PutTape<Self::Tape, Output = Self>;
     /// Splits tape off of self
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// # let grads = Gradients::leaky();
     /// let a: Tensor<Rank1<5>, f32, _, OwnedTape<f32, _>> = dev.zeros().traced(grads);
diff --git a/dfdx-core/src/tensor_ops/abs/mod.rs b/dfdx-core/src/tensor_ops/abs/mod.rs
index 6d4d663dd..361079f34 100644
--- a/dfdx-core/src/tensor_ops/abs/mod.rs
+++ b/dfdx-core/src/tensor_ops/abs/mod.rs
@@ -16,7 +16,7 @@ pub struct AbsKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.abs();
diff --git a/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs b/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs
index 162406fd2..56dda8a04 100644
--- a/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs
+++ b/dfdx-core/src/tensor_ops/accurate_gelu/mod.rs
@@ -28,7 +28,7 @@ pub struct AccurateGeLUKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.accurate_gelu();
diff --git a/dfdx-core/src/tensor_ops/adam/mod.rs b/dfdx-core/src/tensor_ops/adam/mod.rs
index 34ab40b90..fd81e64cc 100644
--- a/dfdx-core/src/tensor_ops/adam/mod.rs
+++ b/dfdx-core/src/tensor_ops/adam/mod.rs
@@ -14,7 +14,7 @@ use super::WeightDecay;
 ///
 /// Changing all default parameters:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// AdamConfig {
 ///     lr: 1e-2,
 ///     betas: [0.1, 0.2],
diff --git a/dfdx-core/src/tensor_ops/add/mod.rs b/dfdx-core/src/tensor_ops/add/mod.rs
index fe7d9cd8c..a70270042 100644
--- a/dfdx-core/src/tensor_ops/add/mod.rs
+++ b/dfdx-core/src/tensor_ops/add/mod.rs
@@ -23,7 +23,7 @@ pub struct ScalarAddKernelOp<E> {
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let r = a + dev.ones();
@@ -32,7 +32,7 @@ pub struct ScalarAddKernelOp<E> {
 ///
 /// Adding a scalar:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let r = a + 1.0;
diff --git a/dfdx-core/src/tensor_ops/boolean/mod.rs b/dfdx-core/src/tensor_ops/boolean/mod.rs
index adf1be2ea..51c9c5def 100644
--- a/dfdx-core/src/tensor_ops/boolean/mod.rs
+++ b/dfdx-core/src/tensor_ops/boolean/mod.rs
@@ -133,7 +133,7 @@ boolean_op_impl!(BitXor, bitxor, xor, scalar_xor);
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, true, false]);
 ///
@@ -151,7 +151,7 @@ pub fn bool_not<S: Shape, E: Dtype, D: Device<E>>(inp: &Tensor<S, bool, D>) -> T
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, false, true, true]);
 /// let b = dev.tensor([false, true, false, true]);
@@ -165,7 +165,7 @@ pub fn bool_not<S: Shape, E: Dtype, D: Device<E>>(inp: &Tensor<S, bool, D>) -> T
 ///
 /// And-ing with a scalar:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, true, false]);
 ///
@@ -186,7 +186,7 @@ pub fn bool_and<S: Shape, E: Dtype, D: Device<E>>(
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, false, true, true]);
 /// let b = dev.tensor([false, true, false, true]);
@@ -200,7 +200,7 @@ pub fn bool_and<S: Shape, E: Dtype, D: Device<E>>(
 ///
 /// And-ing with a scalar:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, true, false]);
 ///
@@ -221,7 +221,7 @@ pub fn bool_or<S: Shape, E: Dtype, D: Device<E>>(
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, false, true, true]);
 /// let b = dev.tensor([false, true, false, true]);
@@ -235,7 +235,7 @@ pub fn bool_or<S: Shape, E: Dtype, D: Device<E>>(
 ///
 /// And-ing with a scalar:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([false, true, false]);
 ///
diff --git a/dfdx-core/src/tensor_ops/broadcast_to.rs b/dfdx-core/src/tensor_ops/broadcast_to.rs
index 7b0bbebfd..90dfa93a9 100644
--- a/dfdx-core/src/tensor_ops/broadcast_to.rs
+++ b/dfdx-core/src/tensor_ops/broadcast_to.rs
@@ -6,7 +6,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Use shape generic or output type to dictate what shape you want:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 7>, f32, _> = dev.zeros();
 /// // broadcast axis 1
@@ -17,7 +17,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Use axes generic to dis-ambiguate:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank1<1>, f32, _> = dev.zeros();
 /// // It's ambiguous what axes to broadcast here - explicitly say axes 0 and 2
diff --git a/dfdx-core/src/tensor_ops/choose/mod.rs b/dfdx-core/src/tensor_ops/choose/mod.rs
index 9f011dd4e..a1861e448 100644
--- a/dfdx-core/src/tensor_ops/choose/mod.rs
+++ b/dfdx-core/src/tensor_ops/choose/mod.rs
@@ -30,7 +30,7 @@ pub trait ChooseKernel<E: Dtype>: Storage<E> + Storage<bool> {
 /// Choose values from two tensors using a boolean mask. Equivalent to `torch.where` from pytorch.
 ///
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let cond: Tensor<Rank1<3>, bool, _> = dev.tensor([true, false, true]);
 /// let a: Tensor<Rank1<3>, f32, _> = dev.tensor([1.0, 2.0, 3.0]);
diff --git a/dfdx-core/src/tensor_ops/clamp/mod.rs b/dfdx-core/src/tensor_ops/clamp/mod.rs
index b58b624aa..b23f6a035 100644
--- a/dfdx-core/src/tensor_ops/clamp/mod.rs
+++ b/dfdx-core/src/tensor_ops/clamp/mod.rs
@@ -17,7 +17,7 @@ pub struct ClampKernelOp<E> {
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, -0.5, 0.0, 0.5, 1.0]);
 /// let r = t.clamp(-0.5, 0.5);
diff --git a/dfdx-core/src/tensor_ops/cmp/mod.rs b/dfdx-core/src/tensor_ops/cmp/mod.rs
index b499ef559..cd1c3cc08 100644
--- a/dfdx-core/src/tensor_ops/cmp/mod.rs
+++ b/dfdx-core/src/tensor_ops/cmp/mod.rs
@@ -49,7 +49,7 @@ pub enum LeKernelOp {}
 ///
 /// Examples:
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.2, 3.4, -0.5, -6.7]);
 /// let b = dev.tensor([1.2, 0.0, 3.14, -6.7]);
@@ -57,7 +57,7 @@ pub enum LeKernelOp {}
 /// assert_eq!(r.array(), [true, false, false, true]);
 /// ```
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2, -1, 0, 1, 2]);
 /// let r = a.eq(1);
@@ -74,7 +74,7 @@ pub fn eq<S: Shape, E, D: CmpKernel<EqKernelOp, E>, T: Tape<E, D>>(
 ///
 /// Examples:
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.2, 3.4, -0.5, -6.7]);
 /// let b = dev.tensor([1.2, 0.0, 3.14, -6.7]);
@@ -82,7 +82,7 @@ pub fn eq<S: Shape, E, D: CmpKernel<EqKernelOp, E>, T: Tape<E, D>>(
 /// assert_eq!(r.array(), [false, true, true, false]);
 /// ```
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2, -1, 0, 1, 2]);
 /// let r = a.ne(1);
@@ -99,7 +99,7 @@ pub fn ne<S: Shape, E, D: CmpKernel<NeKernelOp, E>, T: Tape<E, D>>(
 ///
 /// Examples:
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.2, 3.4, -0.5, -1.0]);
 /// let b = dev.tensor([1.2, 0.0, 3.14, -6.7]);
@@ -107,7 +107,7 @@ pub fn ne<S: Shape, E, D: CmpKernel<NeKernelOp, E>, T: Tape<E, D>>(
 /// assert_eq!(r.array(), [false, true, false, true]);
 /// ```
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2, -1, 0, 1, 2]);
 /// let r = a.gt(-1);
@@ -124,7 +124,7 @@ pub fn gt<S: Shape, E, D: CmpKernel<GtKernelOp, E>, T: Tape<E, D>>(
 ///
 /// Examples:
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.2, 3.4, -0.5, -1.0]);
 /// let b = dev.tensor([1.2, 0.0, 3.14, -6.7]);
@@ -132,7 +132,7 @@ pub fn gt<S: Shape, E, D: CmpKernel<GtKernelOp, E>, T: Tape<E, D>>(
 /// assert_eq!(r.array(), [true, true, false, true]);
 /// ```
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2, -1, 0, 1, 2]);
 /// let r = a.ge(-1);
@@ -149,7 +149,7 @@ pub fn ge<S: Shape, E, D: CmpKernel<GeKernelOp, E>, T: Tape<E, D>>(
 ///
 /// Examples:
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.2, 3.4, -0.5, -1.0]);
 /// let b = dev.tensor([1.2, 0.0, 3.14, -6.7]);
@@ -157,7 +157,7 @@ pub fn ge<S: Shape, E, D: CmpKernel<GeKernelOp, E>, T: Tape<E, D>>(
 /// assert_eq!(r.array(), [false, false, true, false]);
 /// ```
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2, -1, 0, 1, 2]);
 /// let r = a.lt(1);
@@ -174,7 +174,7 @@ pub fn lt<S: Shape, E, D: CmpKernel<LtKernelOp, E>, T: Tape<E, D>>(
 ///
 /// Examples:
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.2, 3.4, -0.5, -1.0]);
 /// let b = dev.tensor([1.2, 0.0, 3.14, -6.7]);
@@ -182,7 +182,7 @@ pub fn lt<S: Shape, E, D: CmpKernel<LtKernelOp, E>, T: Tape<E, D>>(
 /// assert_eq!(r.array(), [true, false, true, false]);
 /// ```
 /// ```
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2, -1, 0, 1, 2]);
 /// let r = a.le(1);
diff --git a/dfdx-core/src/tensor_ops/concat/mod.rs b/dfdx-core/src/tensor_ops/concat/mod.rs
index d75f5ea89..0ad3a9669 100644
--- a/dfdx-core/src/tensor_ops/concat/mod.rs
+++ b/dfdx-core/src/tensor_ops/concat/mod.rs
@@ -10,7 +10,7 @@ mod cuda_kernel;
 ///
 /// Concat with const dims **requires nightly**:
 /// ```ignore
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
 /// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
@@ -19,7 +19,7 @@ mod cuda_kernel;
 ///
 /// Concat with usize dims:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(2, Const));
 /// let b: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(4, Const));
diff --git a/dfdx-core/src/tensor_ops/concat_along/mod.rs b/dfdx-core/src/tensor_ops/concat_along/mod.rs
index 71ab0291e..28282419f 100644
--- a/dfdx-core/src/tensor_ops/concat_along/mod.rs
+++ b/dfdx-core/src/tensor_ops/concat_along/mod.rs
@@ -12,7 +12,7 @@ mod cuda_kernel;
 ///
 /// Along Axis 0:
 /// ```ignore
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
 /// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
@@ -21,7 +21,7 @@ mod cuda_kernel;
 ///
 /// Along Axis 1:
 /// ```ignore
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
 /// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
@@ -31,7 +31,7 @@ mod cuda_kernel;
 /// # [usize] dims
 /// Along Axis 0:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(2, Const));
 /// let b: Tensor<(usize, Const<3>), f32, _> = dev.zeros_like(&(4, Const));
@@ -40,7 +40,7 @@ mod cuda_kernel;
 ///
 /// Along Axis 1:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<(Const<2>, usize), f32, _> = dev.zeros_like(&(Const, 2));
 /// let b: Tensor<(Const<2>, usize), f32, _> = dev.zeros_like(&(Const, 4));
diff --git a/dfdx-core/src/tensor_ops/conv1d/mod.rs b/dfdx-core/src/tensor_ops/conv1d/mod.rs
index ad4335d99..25bdddbee 100644
--- a/dfdx-core/src/tensor_ops/conv1d/mod.rs
+++ b/dfdx-core/src/tensor_ops/conv1d/mod.rs
@@ -51,7 +51,7 @@ pub(super) trait Conv1DKernel<E: Dtype>: Storage<E> {
 /// [Const] dims **require nightly**:
 /// ```ignore
 /// #![feature(generic_const_exprs)]
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank4<2, 3, 32>, f32, _> = dev.sample_normal();
 /// let w: Tensor<Rank4<6, 3, 3>, f32, _> = dev.sample_normal();
@@ -65,7 +65,7 @@ pub(super) trait Conv1DKernel<E: Dtype>: Storage<E> {
 ///
 /// [usize] dims can be used on stable:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<_, f32, _> = dev.sample_normal_like(&(
 ///     2,  // batch size
diff --git a/dfdx-core/src/tensor_ops/conv2d/mod.rs b/dfdx-core/src/tensor_ops/conv2d/mod.rs
index eef6ff179..c61b4cb7d 100644
--- a/dfdx-core/src/tensor_ops/conv2d/mod.rs
+++ b/dfdx-core/src/tensor_ops/conv2d/mod.rs
@@ -57,7 +57,7 @@ pub(super) trait Conv2DKernel<E: Dtype>: Storage<E> {
 /// [Const] dims **require nightly**:
 /// ```ignore
 /// #![feature(generic_const_exprs)]
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank4<2, 3, 32, 32>, f32, _> = dev.sample_normal();
 /// let w: Tensor<Rank4<6, 3, 3, 3>, f32, _> = dev.sample_normal();
@@ -71,7 +71,7 @@ pub(super) trait Conv2DKernel<E: Dtype>: Storage<E> {
 ///
 /// [usize] dims can be used on stable:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<_, f32, _> = dev.sample_normal_like(&(
 ///     2,  // batch size
diff --git a/dfdx-core/src/tensor_ops/cos/mod.rs b/dfdx-core/src/tensor_ops/cos/mod.rs
index 1a6ea4a53..01bd69d4c 100644
--- a/dfdx-core/src/tensor_ops/cos/mod.rs
+++ b/dfdx-core/src/tensor_ops/cos/mod.rs
@@ -16,7 +16,7 @@ pub struct CosKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.cos();
diff --git a/dfdx-core/src/tensor_ops/div/mod.rs b/dfdx-core/src/tensor_ops/div/mod.rs
index 6bcf93fd6..1907b3971 100644
--- a/dfdx-core/src/tensor_ops/div/mod.rs
+++ b/dfdx-core/src/tensor_ops/div/mod.rs
@@ -20,7 +20,7 @@ pub struct BinaryDivKernelOp;
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let b = dev.tensor([[1.0, 0.5, 1.0], [0.5, 1.0, 3.0]]);
@@ -30,7 +30,7 @@ pub struct BinaryDivKernelOp;
 ///
 /// Scalar example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let r = a / 2.0;
diff --git a/dfdx-core/src/tensor_ops/dropout/mod.rs b/dfdx-core/src/tensor_ops/dropout/mod.rs
index 73d5fe907..1528853db 100644
--- a/dfdx-core/src/tensor_ops/dropout/mod.rs
+++ b/dfdx-core/src/tensor_ops/dropout/mod.rs
@@ -36,7 +36,7 @@ pub trait DropoutKernel<E: Dtype>: Storage<E> + RandomU64 {
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([1.0f32, 2.0, 3.0, 4.0]);
 /// let r = t.dropout(0.5);
diff --git a/dfdx-core/src/tensor_ops/exp/mod.rs b/dfdx-core/src/tensor_ops/exp/mod.rs
index e8c7abe70..d17c56f72 100644
--- a/dfdx-core/src/tensor_ops/exp/mod.rs
+++ b/dfdx-core/src/tensor_ops/exp/mod.rs
@@ -16,7 +16,7 @@ pub struct ExpKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.exp();
diff --git a/dfdx-core/src/tensor_ops/fast_gelu/mod.rs b/dfdx-core/src/tensor_ops/fast_gelu/mod.rs
index 02a1de175..831ef5db6 100644
--- a/dfdx-core/src/tensor_ops/fast_gelu/mod.rs
+++ b/dfdx-core/src/tensor_ops/fast_gelu/mod.rs
@@ -23,7 +23,7 @@ pub struct FastGeLUKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.gelu();
diff --git a/dfdx-core/src/tensor_ops/huber_error/mod.rs b/dfdx-core/src/tensor_ops/huber_error/mod.rs
index 335c28e29..b2ee13c39 100644
--- a/dfdx-core/src/tensor_ops/huber_error/mod.rs
+++ b/dfdx-core/src/tensor_ops/huber_error/mod.rs
@@ -21,7 +21,7 @@ pub struct HuberErrorKernelOp<E> {
 /// 2. otherwise: `delta * (|x - y| - 0.5 * delta)`
 ///
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1.0, 1.0, 1.0]);
 /// let b = dev.tensor([1.5, 1.75, 2.5]);
diff --git a/dfdx-core/src/tensor_ops/ln/mod.rs b/dfdx-core/src/tensor_ops/ln/mod.rs
index 93f8d9bd4..a144e3b6f 100644
--- a/dfdx-core/src/tensor_ops/ln/mod.rs
+++ b/dfdx-core/src/tensor_ops/ln/mod.rs
@@ -16,7 +16,7 @@ pub struct LnKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.ln();
diff --git a/dfdx-core/src/tensor_ops/log_softmax.rs b/dfdx-core/src/tensor_ops/log_softmax.rs
index d6cfe1631..9a8023899 100644
--- a/dfdx-core/src/tensor_ops/log_softmax.rs
+++ b/dfdx-core/src/tensor_ops/log_softmax.rs
@@ -7,7 +7,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank3<2, 3, 5>, f32, _> = dev.zeros();
 /// let _ = t.log_softmax::<Axis<2>>();
@@ -15,7 +15,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Using multi axis log_softmax:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// # let t: Tensor<Rank3<2, 3, 5>, f32, _> = dev.zeros();
 /// let _ = t.log_softmax::<Axes2<0, 2>>();
diff --git a/dfdx-core/src/tensor_ops/logsumexp_to.rs b/dfdx-core/src/tensor_ops/logsumexp_to.rs
index a188070c8..8f43abf4f 100644
--- a/dfdx-core/src/tensor_ops/logsumexp_to.rs
+++ b/dfdx-core/src/tensor_ops/logsumexp_to.rs
@@ -11,7 +11,7 @@ pub trait LogSumExpTo: HasErr + HasShape {
     ///
     /// Example:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank3<2, 4, 6>, f32, _> = dev.zeros();
     /// let _ = t.logsumexp::<Rank2<2, 4>, _>(); // or `logsumexp::<_, Axis<2>>()`
@@ -19,7 +19,7 @@ pub trait LogSumExpTo: HasErr + HasShape {
     ///
     /// Multi axis logsumexp:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// # let t: Tensor<Rank3<2, 4, 6>, f32, _> = dev.zeros();
     /// let _ = t.logsumexp::<Rank1<4>, _>(); // or `logsumexp::<_, Axes2<0, 2>>()`
diff --git a/dfdx-core/src/tensor_ops/matmul/mod.rs b/dfdx-core/src/tensor_ops/matmul/mod.rs
index 65b71bb4d..5247f4f88 100644
--- a/dfdx-core/src/tensor_ops/matmul/mod.rs
+++ b/dfdx-core/src/tensor_ops/matmul/mod.rs
@@ -17,7 +17,7 @@ use super::reshape_to::{ReshapeKernel, ReshapeTo};
 /// # Examples
 /// 1. Matrix & Matrix
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank2<3, 10>, f32, _> = dev.zeros();
 /// let y: Tensor<Rank2<10, 5>, f32, _> = dev.zeros();
@@ -26,7 +26,7 @@ use super::reshape_to::{ReshapeKernel, ReshapeTo};
 ///
 /// 2. Vector x Matrix
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank1<2>, f32, _> = dev.zeros();
 /// let y: Tensor<Rank2<2, 4>, f32, _> = dev.zeros();
@@ -35,7 +35,7 @@ use super::reshape_to::{ReshapeKernel, ReshapeTo};
 ///
 /// 3. Vector x Vector
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank1<2>, f32, _> = dev.zeros();
 /// let y: Tensor<Rank1<4>, f32, _> = dev.zeros();
@@ -44,7 +44,7 @@ use super::reshape_to::{ReshapeKernel, ReshapeTo};
 ///
 /// 4. Batched matmul
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank3<10, 3, 2>, f32, _> = dev.zeros();
 /// let y: Tensor<Rank3<10, 2, 4>, f32, _> = dev.zeros();
@@ -53,7 +53,7 @@ use super::reshape_to::{ReshapeKernel, ReshapeTo};
 ///
 /// 5. Broadcasted matmul
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let x: Tensor<Rank3<10, 3, 2>, f32, _> = dev.zeros();
 /// let y: Tensor<Rank2<2, 4>, f32, _> = dev.zeros();
@@ -179,7 +179,7 @@ where
 {
     type Output = Tensor<(M, N), E, D, T>;
     /// ```compile_fail
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let x: Tensor<Rank2<3, 2>, f32, _> = dev.zeros();
     /// let y: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
@@ -216,7 +216,7 @@ where
 {
     type Output = Tensor<(B, M, N), E, D, T>;
     /// ```compile_fail
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let x: Tensor<Rank3<1, 3, 2>, f32, _> = dev.zeros();
     /// let y: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
@@ -254,7 +254,7 @@ where
 {
     type Output = Tensor<(B, M, N), E, D, T>;
     /// ```compile_fail
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let x: Tensor<Rank3<1, 3, 2>, f32, _> = dev.zeros();
     /// let y: Tensor<Rank3<1, 3, 4>, f32, _> = dev.zeros();
@@ -293,7 +293,7 @@ where
 {
     type Output = Tensor<(B, S, M, N), E, D, T>;
     /// ```compile_fail
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let x: Tensor<Rank4<1, 5, 3, 2>, f32, _> = dev.zeros();
     /// let y: Tensor<Rank4<1, 5, 3, 4>, f32, _> = dev.zeros();
diff --git a/dfdx-core/src/tensor_ops/max_to/mod.rs b/dfdx-core/src/tensor_ops/max_to/mod.rs
index 7a39b6859..93eac94b3 100644
--- a/dfdx-core/src/tensor_ops/max_to/mod.rs
+++ b/dfdx-core/src/tensor_ops/max_to/mod.rs
@@ -33,7 +33,7 @@ pub trait MaxTo: HasErr + HasShape {
     ///
     /// Example reducing a single axis:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
     /// let r = t.max::<Rank1<2>, _>(); // or `max::<_, Axis<1>>()`
@@ -42,7 +42,7 @@ pub trait MaxTo: HasErr + HasShape {
     ///
     /// Reducing multiple axes:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// # let t = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
     /// let r = t.max::<Rank0, _>();
diff --git a/dfdx-core/src/tensor_ops/maximum/mod.rs b/dfdx-core/src/tensor_ops/maximum/mod.rs
index 5b514ddc8..4d042c468 100644
--- a/dfdx-core/src/tensor_ops/maximum/mod.rs
+++ b/dfdx-core/src/tensor_ops/maximum/mod.rs
@@ -16,7 +16,7 @@ pub struct MaximumKernelOp;
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let b = dev.tensor([[1.0, 0.5, 1.0], [-2.0, 2.0, -3.5]]);
diff --git a/dfdx-core/src/tensor_ops/mean_to.rs b/dfdx-core/src/tensor_ops/mean_to.rs
index c87632cae..fb0b2f851 100644
--- a/dfdx-core/src/tensor_ops/mean_to.rs
+++ b/dfdx-core/src/tensor_ops/mean_to.rs
@@ -7,7 +7,7 @@ pub trait MeanTo: HasErr + HasShape {
     ///
     /// Example:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
     /// let r = t.mean::<Rank0, _>(); // or `mean::<_, Axes2<0, 1>>()`
@@ -16,7 +16,7 @@ pub trait MeanTo: HasErr + HasShape {
     ///
     /// Reducing 1 axis:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
     /// let r = t.mean::<Rank1<2>, _>(); // or `mean::<_, Axis<1>>()`
diff --git a/dfdx-core/src/tensor_ops/min_to/mod.rs b/dfdx-core/src/tensor_ops/min_to/mod.rs
index c534c583f..ae426f4f1 100644
--- a/dfdx-core/src/tensor_ops/min_to/mod.rs
+++ b/dfdx-core/src/tensor_ops/min_to/mod.rs
@@ -33,7 +33,7 @@ pub trait MinTo: HasErr + HasShape {
     ///
     /// Example reducing a single axis:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
     /// let r = t.min::<Rank1<2>, _>(); // or `min::<_, Axis<1>>()`
@@ -42,7 +42,7 @@ pub trait MinTo: HasErr + HasShape {
     ///
     /// Reducing multiple axes:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// # let t = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
     /// let r = t.min::<Rank0, _>();
diff --git a/dfdx-core/src/tensor_ops/minimum/mod.rs b/dfdx-core/src/tensor_ops/minimum/mod.rs
index c0d766082..e93c38426 100644
--- a/dfdx-core/src/tensor_ops/minimum/mod.rs
+++ b/dfdx-core/src/tensor_ops/minimum/mod.rs
@@ -16,7 +16,7 @@ pub struct MinimumKernelOp;
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let b = dev.tensor([[1.0, 0.5, 1.0], [-2.0, 2.0, -3.5]]);
diff --git a/dfdx-core/src/tensor_ops/mod.rs b/dfdx-core/src/tensor_ops/mod.rs
index 2d1ce188c..c51040eeb 100644
--- a/dfdx-core/src/tensor_ops/mod.rs
+++ b/dfdx-core/src/tensor_ops/mod.rs
@@ -43,7 +43,7 @@
 //!
 //! For example:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t: Tensor<Rank3<2, 4, 6>, f32, _> = dev.zeros();
 //! // shape version
@@ -74,7 +74,7 @@
 //!
 //! For example:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t: Tensor<Rank1<4>, f32, _> = dev.zeros();
 //! // shape version
@@ -85,7 +85,7 @@
 //!
 //! Rust can also infer the output type if you use it in another operation:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let big: Tensor<Rank2<2, 5>, f32, _> = dev.zeros();
 //! let small: Tensor<Rank1<5>, f32, _> = dev.zeros();
@@ -96,7 +96,7 @@
 //!
 //! Permuting has an identical interface to broadcasts/reductions:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t: Tensor<Rank3<2, 3, 4>, f32, _> = dev.zeros();
 //! // shape version
@@ -113,7 +113,7 @@
 //!
 //! For example you can select from the 0th axis like so:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
 //! let r: Tensor<Rank1<3>, f32, _> = t.select(dev.tensor(1));
@@ -122,7 +122,7 @@
 //!
 //! Or you can gather from the 0th axis to select multiple entries:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
 //! let r: Tensor<Rank2<3, 3>, f32, _> = t.gather(dev.tensor([1, 1, 0]));
@@ -139,7 +139,7 @@
 //! But you can use [BroadcastTo] to make this easy! In this example we select
 //! the same index from the 1st axis of a tensor:
 //! ```rust
-//! # use dfdx::prelude::*;
+//! # use dfdx_core::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let t = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
 //! let r = t.select::<Rank1<2>, _>(dev.tensor(1).broadcast());
diff --git a/dfdx-core/src/tensor_ops/mul/mod.rs b/dfdx-core/src/tensor_ops/mul/mod.rs
index 4b66bfbcb..0d345de14 100644
--- a/dfdx-core/src/tensor_ops/mul/mod.rs
+++ b/dfdx-core/src/tensor_ops/mul/mod.rs
@@ -20,7 +20,7 @@ pub struct ScalarMulKernelOp<E> {
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let r = a * dev.ones();
@@ -29,7 +29,7 @@ pub struct ScalarMulKernelOp<E> {
 ///
 /// Scalar example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let r = a * 2.0;
diff --git a/dfdx-core/src/tensor_ops/nans_to/mod.rs b/dfdx-core/src/tensor_ops/nans_to/mod.rs
index 2909238df..9b511d5e7 100644
--- a/dfdx-core/src/tensor_ops/nans_to/mod.rs
+++ b/dfdx-core/src/tensor_ops/nans_to/mod.rs
@@ -16,7 +16,7 @@ pub struct NansToKernelOp<E>(E);
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([1.0, f32::NAN, f32::NAN, 4.0]);
 /// let r = t.nans_to(0.0);
diff --git a/dfdx-core/src/tensor_ops/negate/mod.rs b/dfdx-core/src/tensor_ops/negate/mod.rs
index bc36bb95b..89bf985c6 100644
--- a/dfdx-core/src/tensor_ops/negate/mod.rs
+++ b/dfdx-core/src/tensor_ops/negate/mod.rs
@@ -14,7 +14,7 @@ pub struct NegateKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([-2.0, 0.0, 5.0]);
 /// let r = -a;
diff --git a/dfdx-core/src/tensor_ops/normalize.rs b/dfdx-core/src/tensor_ops/normalize.rs
index 019e705a7..49e13994d 100644
--- a/dfdx-core/src/tensor_ops/normalize.rs
+++ b/dfdx-core/src/tensor_ops/normalize.rs
@@ -10,7 +10,7 @@ use super::{BroadcastTo, Device, MeanTo, TryAdd, TryDiv, TrySub};
 ///
 /// Normalizing a single axis:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
 /// let _ = t.normalize::<Axis<1>>(1e-5);
diff --git a/dfdx-core/src/tensor_ops/permute_to.rs b/dfdx-core/src/tensor_ops/permute_to.rs
index ad07a63a3..920d46e62 100644
--- a/dfdx-core/src/tensor_ops/permute_to.rs
+++ b/dfdx-core/src/tensor_ops/permute_to.rs
@@ -6,7 +6,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Option 1: Specifying shape generic:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
 /// let b: Tensor<Rank2<3, 2>, f32, _> = a.permute::<Rank2<3, 2>, _>();
@@ -15,7 +15,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Option 2: Specifying axes generic:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);
 /// let b: Tensor<Rank2<3, 2>, f32, _> = a.permute::<_, Axes2<1, 0>>();
diff --git a/dfdx-core/src/tensor_ops/pow/mod.rs b/dfdx-core/src/tensor_ops/pow/mod.rs
index ecef22110..1d33ad243 100644
--- a/dfdx-core/src/tensor_ops/pow/mod.rs
+++ b/dfdx-core/src/tensor_ops/pow/mod.rs
@@ -16,7 +16,7 @@ pub struct PowfKernelOp<E>(E);
 
 /// Raises to a float power; `t^i`.
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.powf(-3.2);
@@ -42,7 +42,7 @@ impl<S: Shape, E: Dtype, D: UnaryKernel<PowfKernelOp<E>, E>, T: Tape<E, D>> Tens
 
 /// Raises to an integer power; `t^i`.
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.powi(3);
diff --git a/dfdx-core/src/tensor_ops/prelu.rs b/dfdx-core/src/tensor_ops/prelu.rs
index b6cc38ac9..0bc018ca5 100644
--- a/dfdx-core/src/tensor_ops/prelu.rs
+++ b/dfdx-core/src/tensor_ops/prelu.rs
@@ -11,7 +11,7 @@ use super::{cmp::*, BroadcastTo, ChooseFrom, Device, TryMul};
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let a = dev.tensor([0.05,0.05,0.05,0.05]);
@@ -42,7 +42,7 @@ pub fn leakyrelu<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>>(
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let a = dev.tensor([0.05,0.05,0.05,0.05]);
diff --git a/dfdx-core/src/tensor_ops/realize_to.rs b/dfdx-core/src/tensor_ops/realize_to.rs
index e29fd91be..455ed62f1 100644
--- a/dfdx-core/src/tensor_ops/realize_to.rs
+++ b/dfdx-core/src/tensor_ops/realize_to.rs
@@ -3,7 +3,7 @@ use crate::{shapes::*, tensor::*};
 /// Realizes the concrete shape of the tensor as another compatable shape,
 /// or returns the original tensor if the new shape's dimensions are incompatable.
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<2, 3>, f32, _> = dev.zeros();
 /// let a = a.realize::<(usize, usize)>();
diff --git a/dfdx-core/src/tensor_ops/recip/mod.rs b/dfdx-core/src/tensor_ops/recip/mod.rs
index 78eb28792..3a1d9157d 100644
--- a/dfdx-core/src/tensor_ops/recip/mod.rs
+++ b/dfdx-core/src/tensor_ops/recip/mod.rs
@@ -14,7 +14,7 @@ pub struct RecipKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.recip();
diff --git a/dfdx-core/src/tensor_ops/relu/mod.rs b/dfdx-core/src/tensor_ops/relu/mod.rs
index c1ac134db..0701741c3 100644
--- a/dfdx-core/src/tensor_ops/relu/mod.rs
+++ b/dfdx-core/src/tensor_ops/relu/mod.rs
@@ -16,7 +16,7 @@ pub struct ReLUKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.relu();
diff --git a/dfdx-core/src/tensor_ops/reshape_to/mod.rs b/dfdx-core/src/tensor_ops/reshape_to/mod.rs
index 35f4ed781..04f058313 100644
--- a/dfdx-core/src/tensor_ops/reshape_to/mod.rs
+++ b/dfdx-core/src/tensor_ops/reshape_to/mod.rs
@@ -26,7 +26,7 @@ pub trait ReshapeKernel<E: Dtype>: Storage<E> {
 ///
 /// Compile time reshapes:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank2<2, 4>, f32, _> = dev.zeros();
 /// let t: Tensor<Rank1<8>, f32, _> = t.reshape();
@@ -34,7 +34,7 @@ pub trait ReshapeKernel<E: Dtype>: Storage<E> {
 ///
 /// Compile time failure:
 /// ```compile_fail
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank2<2, 4>, f32, _> = dev.zeros();
 /// let t: Tensor<Rank1<7>, f32, _> = t.reshape();
@@ -42,7 +42,7 @@ pub trait ReshapeKernel<E: Dtype>: Storage<E> {
 ///
 /// Runtime reshapes:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank2<2, 4>, f32, _> = dev.zeros();
 /// let t: Tensor<(usize, ), f32, _> = t.reshape_like(&(8, ));
diff --git a/dfdx-core/src/tensor_ops/roll/mod.rs b/dfdx-core/src/tensor_ops/roll/mod.rs
index 933d3b259..ad7318792 100644
--- a/dfdx-core/src/tensor_ops/roll/mod.rs
+++ b/dfdx-core/src/tensor_ops/roll/mod.rs
@@ -32,7 +32,7 @@ pub trait RollKernel<E: Dtype>: Storage<E> {
 /// Shifts data along an axis by a specified amount.
 ///
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([1.0, 2.0, 3.0, 4.0]);
 /// let r = t.roll::<Axis<0>>(1);
@@ -41,7 +41,7 @@ pub trait RollKernel<E: Dtype>: Storage<E> {
 ///
 /// Won't compile if you try to roll an axis that doesn't exist:
 /// ```compile_fail
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([1.0, 2.0, 3.0, 4.0]);
 /// let r = t.roll::<Axis<3>>(1);
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 58558ebe7..f6178e9c7 100644
--- a/dfdx-core/src/tensor_ops/select_and_gather/mod.rs
+++ b/dfdx-core/src/tensor_ops/select_and_gather/mod.rs
@@ -61,7 +61,7 @@ pub trait RemoveDimKernel<E: Dtype>: Storage<E> + Storage<usize> {
 ///
 /// Here is an example selecting from a 2d tensor:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 5>, f32, _> = dev.zeros();
 ///
@@ -134,7 +134,7 @@ impl<Src: Shape, E: Dtype, D: RemoveDimKernel<E>, T: Tape<E, D>> SelectTo<E, D>
 ///
 /// Here is an example gathering from a 2d tensor:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 5>, f32, _> = dev.zeros();
 ///
diff --git a/dfdx-core/src/tensor_ops/sgd/mod.rs b/dfdx-core/src/tensor_ops/sgd/mod.rs
index 6b060feb7..3515f9f57 100644
--- a/dfdx-core/src/tensor_ops/sgd/mod.rs
+++ b/dfdx-core/src/tensor_ops/sgd/mod.rs
@@ -14,7 +14,7 @@ use super::optim::{Momentum, WeightDecay};
 ///
 /// Using different learning rate:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// SgdConfig {
 ///     lr: 1e-1,
 ///     momentum: None,
@@ -24,7 +24,7 @@ use super::optim::{Momentum, WeightDecay};
 ///
 /// Using classic momentum:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// SgdConfig {
 ///     lr: 1e-2,
 ///     momentum: Some(Momentum::Classic(0.5)),
@@ -34,7 +34,7 @@ use super::optim::{Momentum, WeightDecay};
 ///
 /// Using nesterov momentum:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// SgdConfig {
 ///     lr: 1e-3,
 ///     momentum: Some(Momentum::Nesterov(0.25)),
@@ -44,7 +44,7 @@ use super::optim::{Momentum, WeightDecay};
 ///
 /// Using L2 weight decay:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// SgdConfig {
 ///     lr: 1e-3,
 ///     momentum: None,
@@ -54,7 +54,7 @@ use super::optim::{Momentum, WeightDecay};
 ///
 /// Using decoupled weight decay:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// SgdConfig {
 ///     lr: 1e-3,
 ///     momentum: None,
diff --git a/dfdx-core/src/tensor_ops/sigmoid/mod.rs b/dfdx-core/src/tensor_ops/sigmoid/mod.rs
index 6e3a0661c..316b85564 100644
--- a/dfdx-core/src/tensor_ops/sigmoid/mod.rs
+++ b/dfdx-core/src/tensor_ops/sigmoid/mod.rs
@@ -16,7 +16,7 @@ pub struct SigmoidKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.sigmoid();
diff --git a/dfdx-core/src/tensor_ops/sin/mod.rs b/dfdx-core/src/tensor_ops/sin/mod.rs
index 5a3fe2f05..841d1ebba 100644
--- a/dfdx-core/src/tensor_ops/sin/mod.rs
+++ b/dfdx-core/src/tensor_ops/sin/mod.rs
@@ -16,7 +16,7 @@ pub struct SinKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.sin();
diff --git a/dfdx-core/src/tensor_ops/slice/mod.rs b/dfdx-core/src/tensor_ops/slice/mod.rs
index 8e1d60117..9dce18815 100644
--- a/dfdx-core/src/tensor_ops/slice/mod.rs
+++ b/dfdx-core/src/tensor_ops/slice/mod.rs
@@ -29,7 +29,7 @@ pub trait SliceKernel<E: Unit>: Storage<E> {
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev = Cpu::default();
 /// let a = dev.tensor([
 ///     [1., 2.],
diff --git a/dfdx-core/src/tensor_ops/softmax.rs b/dfdx-core/src/tensor_ops/softmax.rs
index 667c8e6bb..4390319da 100644
--- a/dfdx-core/src/tensor_ops/softmax.rs
+++ b/dfdx-core/src/tensor_ops/softmax.rs
@@ -10,7 +10,7 @@ use crate::{shapes::*, tensor::*};
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank3<2, 3, 5>, f32, _> = dev.zeros();
 /// let _ = t.softmax::<Axis<2>>();
diff --git a/dfdx-core/src/tensor_ops/sqrt/mod.rs b/dfdx-core/src/tensor_ops/sqrt/mod.rs
index bbe03cb9a..6fa472f1c 100644
--- a/dfdx-core/src/tensor_ops/sqrt/mod.rs
+++ b/dfdx-core/src/tensor_ops/sqrt/mod.rs
@@ -16,7 +16,7 @@ pub struct SqrtKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.sqrt();
diff --git a/dfdx-core/src/tensor_ops/square/mod.rs b/dfdx-core/src/tensor_ops/square/mod.rs
index d7361b12b..aadab4f17 100644
--- a/dfdx-core/src/tensor_ops/square/mod.rs
+++ b/dfdx-core/src/tensor_ops/square/mod.rs
@@ -16,7 +16,7 @@ pub struct SquareKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.square();
diff --git a/dfdx-core/src/tensor_ops/stack/mod.rs b/dfdx-core/src/tensor_ops/stack/mod.rs
index 83bfa4ca8..3a6f8350e 100644
--- a/dfdx-core/src/tensor_ops/stack/mod.rs
+++ b/dfdx-core/src/tensor_ops/stack/mod.rs
@@ -15,7 +15,7 @@ mod cuda_kernel;
 ///
 /// Stacking with an array:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
 /// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
@@ -24,7 +24,7 @@ mod cuda_kernel;
 ///
 /// Stacking with a vec:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
 /// let b: Tensor<Rank2<3, 4>, f32, _> = dev.zeros();
diff --git a/dfdx-core/src/tensor_ops/stddev_to.rs b/dfdx-core/src/tensor_ops/stddev_to.rs
index 0b38796b3..9e116cabb 100644
--- a/dfdx-core/src/tensor_ops/stddev_to.rs
+++ b/dfdx-core/src/tensor_ops/stddev_to.rs
@@ -9,7 +9,7 @@ pub trait StddevTo<E: Dtype>: HasErr + HasShape {
     ///
     /// Examples:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t = dev.tensor([[2.0, 3.0, 4.0], [3.0, 6.0, 9.0]]);
     /// let r = t.stddev::<Rank1<2>, _>(0.0); // or `stddev::<_, Axis<1>>(0.0)`
diff --git a/dfdx-core/src/tensor_ops/sub/mod.rs b/dfdx-core/src/tensor_ops/sub/mod.rs
index af58b74ef..1fdd35813 100644
--- a/dfdx-core/src/tensor_ops/sub/mod.rs
+++ b/dfdx-core/src/tensor_ops/sub/mod.rs
@@ -20,7 +20,7 @@ pub struct ScalarSubKernelOp<E> {
 ///
 /// Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let b = dev.ones();
@@ -30,7 +30,7 @@ pub struct ScalarSubKernelOp<E> {
 ///
 /// Scalar Example:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
 /// let r = a - 1.0;
diff --git a/dfdx-core/src/tensor_ops/sum_to/mod.rs b/dfdx-core/src/tensor_ops/sum_to/mod.rs
index e10939789..d0dae445b 100644
--- a/dfdx-core/src/tensor_ops/sum_to/mod.rs
+++ b/dfdx-core/src/tensor_ops/sum_to/mod.rs
@@ -30,7 +30,7 @@ pub trait SumTo: HasErr + HasShape {
     ///
     /// Example reducing a single axis:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t: Tensor<Rank2<2, 3>, f32, _> = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
     /// let r = t.sum::<Rank1<2>, _>(); // or `sum::<_, Axis<1>>()`
@@ -39,7 +39,7 @@ pub trait SumTo: HasErr + HasShape {
     ///
     /// Reducing multiple axes:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// # let t = dev.tensor([[1.0, 2.0, 3.0], [-1.0, -2.0, -3.0]]);
     /// let r = t.sum::<Rank0, _>(); // or `sum::<_, Axes2<0, 1>>()`
diff --git a/dfdx-core/src/tensor_ops/tanh/mod.rs b/dfdx-core/src/tensor_ops/tanh/mod.rs
index 09c67e53e..7ddf191ba 100644
--- a/dfdx-core/src/tensor_ops/tanh/mod.rs
+++ b/dfdx-core/src/tensor_ops/tanh/mod.rs
@@ -16,7 +16,7 @@ pub struct TanhKernelOp;
 ///
 /// Examples:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t = dev.tensor([-1.0, 0.0, 1.0, 2.0]);
 /// let r = t.tanh();
diff --git a/dfdx-core/src/tensor_ops/to_dtype/mod.rs b/dfdx-core/src/tensor_ops/to_dtype/mod.rs
index 0184613f1..85b0b5ab0 100644
--- a/dfdx-core/src/tensor_ops/to_dtype/mod.rs
+++ b/dfdx-core/src/tensor_ops/to_dtype/mod.rs
@@ -12,7 +12,7 @@ pub trait ToDtypeKernel<E1: Unit, E2: Unit>: Storage<E1> + Storage<E2> {
 ///
 /// Example usage:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let a = dev.tensor([1usize, 2, 3, 4, 5]);
 ///
diff --git a/dfdx-core/src/tensor_ops/upscale2d/mod.rs b/dfdx-core/src/tensor_ops/upscale2d/mod.rs
index 03d9c5e92..5cd849ddc 100644
--- a/dfdx-core/src/tensor_ops/upscale2d/mod.rs
+++ b/dfdx-core/src/tensor_ops/upscale2d/mod.rs
@@ -94,7 +94,7 @@ pub trait GenericUpscale2D<M: UpscaleMethod>: HasErr {
 ///
 /// Compile time upscale:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank3<3, 32, 32>, f32, _> = dev.zeros();
 /// let y: Tensor<Rank3<3, 64, 64>, f32, _> = t.upscale2d(NearestNeighbor);
@@ -102,7 +102,7 @@ pub trait GenericUpscale2D<M: UpscaleMethod>: HasErr {
 ///
 /// Runtime upscale:
 /// ```rust
-/// # use dfdx::prelude::*;
+/// # use dfdx_core::prelude::*;
 /// # let dev: Cpu = Default::default();
 /// let t: Tensor<Rank3<3, 32, 32>, f32, _> = dev.zeros();
 /// let y: Tensor<(Const<3>, usize, usize), f32, _> = t.upscale2d_like(NearestNeighbor, 64, 64);
diff --git a/dfdx-core/src/tensor_ops/var_to.rs b/dfdx-core/src/tensor_ops/var_to.rs
index 8dd5f003d..1a593655c 100644
--- a/dfdx-core/src/tensor_ops/var_to.rs
+++ b/dfdx-core/src/tensor_ops/var_to.rs
@@ -9,7 +9,7 @@ pub trait VarTo: HasErr + HasShape {
     ///
     /// Examples:
     /// ```rust
-    /// # use dfdx::prelude::*;
+    /// # use dfdx_core::prelude::*;
     /// # let dev: Cpu = Default::default();
     /// let t = dev.tensor([[2.0f32, 3.0, 4.0], [3.0, 6.0, 9.0]]);
     /// let r = t.var::<Rank1<2>, _>(); // or `var::<_, Axis<1>>()`
diff --git a/dfdx-derives/src/lib.rs b/dfdx-derives/src/lib.rs
index b0825e62f..f258c679c 100644
--- a/dfdx-derives/src/lib.rs
+++ b/dfdx-derives/src/lib.rs
@@ -25,7 +25,7 @@ macro_rules! has_attr {
 /// Here we have a unit struct that just calls a method on Tensor in the forward:
 ///
 /// ```ignore
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # use dfdx_core::prelude::*;
 /// #[derive(Default, Debug, Clone, Copy, CustomModule)]
 /// pub struct Abs;
@@ -41,7 +41,7 @@ macro_rules! has_attr {
 /// # Using CustomModule on structs with non-parameter fields
 ///
 /// ```ignore
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # use dfdx_core::prelude::*;
 /// #[derive(Default, Debug, Clone, Copy, CustomModule)]
 /// pub struct Reshape<S: Shape>(pub S);
@@ -65,7 +65,7 @@ macro_rules! has_attr {
 /// 3. We must annotate the sub module with `#[module]`
 ///
 /// ```ignore
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # use dfdx_core::prelude::*;
 /// #[derive(Debug, Clone, CustomModule)]
 /// #[built(ResidualMatMul)]
@@ -324,7 +324,7 @@ pub fn custom_module(input: proc_macro::TokenStream) -> proc_macro::TokenStream
 /// 5. linear3
 /// ```ignore
 /// # use dfdx_core::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// #[derive(Debug, Clone, Sequential)]
 /// #[built(Mlp)]
 /// struct MlpConfig {
diff --git a/dfdx/examples/11-module-optim.rs b/dfdx/examples/11-module-optim.rs
index f9b33f05c..6def77c5e 100644
--- a/dfdx/examples/11-module-optim.rs
+++ b/dfdx/examples/11-module-optim.rs
@@ -1,4 +1,4 @@
-//! Intro to dfdx_nn::optim
+//! Intro to dfdx::optim
 
 use dfdx::prelude::*;
 
diff --git a/dfdx/src/lib.rs b/dfdx/src/lib.rs
index a0a94d57e..237ae1ac7 100644
--- a/dfdx/src/lib.rs
+++ b/dfdx/src/lib.rs
@@ -13,7 +13,6 @@
 //! For specifying architecture, you just need the dimensions for the linear, but not the device/dtype:
 //! ```rust
 //! # use dfdx::prelude::*;
-//! # use dfdx_nn::{LinearConfig, LinearConstConfig};
 //! let _: LinearConfig<usize, usize> = LinearConfig::new(3, 5);
 //! let _: LinearConfig<Const<3>, usize> = LinearConfig::new(Const, 5);
 //! let _: LinearConfig<usize, Const<5>> = LinearConfig::new(3, Const);
@@ -30,7 +29,6 @@
 //! We will use [BuildModuleExt::build_module()], an extension trait on devices, to actually construct a model.
 //!
 //! ```rust
-//! # use dfdx_nn::*;
 //! # use dfdx::prelude::*;
 //! let dev: Cpu = Default::default();
 //! let arch = LinearConfig::new(Const::<3>, 5);
@@ -46,7 +44,6 @@
 //! during inference and training.
 //!
 //! ```rust
-//! # use dfdx_nn::*;
 //! # use dfdx::prelude::*;
 //! # let dev: Cpu = Default::default();
 //! let arch = LinearConfig::new(Const::<3>, 5);
@@ -72,7 +69,6 @@
 //!
 //! ```rust
 //! # use dfdx::prelude::*;
-//! # use dfdx_nn::*;
 //! # let dev: Cpu = Default::default();
 //! #[derive(Debug, Clone, Sequential)]
 //! #[built(Mlp)]
@@ -108,7 +104,6 @@
 //!
 //! ```rust
 //! # use dfdx::prelude::*;
-//! # use dfdx_nn::*;
 //! # let dev: Cpu = Default::default();
 //! type Arch = (LinearConstConfig<3, 5>, ReLU, LinearConstConfig<5, 10>);
 //! let mut model = dev.build_module::<f32>(Arch::default());
@@ -134,7 +129,6 @@
 //! a simple example of how to do this:
 //! ```rust
 //! # use dfdx::prelude::*;
-//! # use dfdx_nn::*;
 //! # let dev: Cpu = Default::default();
 //! type Arch = (LinearConstConfig<3, 5>, ReLU, LinearConstConfig<5, 10>);
 //! let arch = Arch::default();
@@ -142,7 +136,7 @@
 //! // 1. allocate gradients for the model
 //! let mut grads = model.alloc_grads();
 //! // 2. create our optimizer
-//! let mut opt = dfdx_nn::optim::Sgd::new(&model, Default::default());
+//! let mut opt = dfdx::nn::optim::Sgd::new(&model, Default::default());
 //! // 3. trace gradients through forward pass
 //! let x: Tensor<Rank2<10, 3>, f32, _> = dev.sample_normal();
 //! let y = model.forward_mut(x.traced(grads));
diff --git a/dfdx/src/nn/layers/add_into.rs b/dfdx/src/nn/layers/add_into.rs
index 885cf82c9..e1d46629f 100644
--- a/dfdx/src/nn/layers/add_into.rs
+++ b/dfdx/src/nn/layers/add_into.rs
@@ -11,7 +11,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = AddInto<(LinearConstConfig<2, 5>, LinearConstConfig<3, 5>)>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/batch_norm1d.rs b/dfdx/src/nn/layers/batch_norm1d.rs
index 2657f44bf..2ae892858 100644
--- a/dfdx/src/nn/layers/batch_norm1d.rs
+++ b/dfdx/src/nn/layers/batch_norm1d.rs
@@ -18,7 +18,7 @@ use crate::prelude::*;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = BatchNorm1DConstConfig<3>;
 /// let bn = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/batch_norm2d.rs b/dfdx/src/nn/layers/batch_norm2d.rs
index 05215f8b1..fcfe97e83 100644
--- a/dfdx/src/nn/layers/batch_norm2d.rs
+++ b/dfdx/src/nn/layers/batch_norm2d.rs
@@ -20,7 +20,7 @@ use crate::prelude::*;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = BatchNorm2DConstConfig<3>;
 /// let bn = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/bias1d.rs b/dfdx/src/nn/layers/bias1d.rs
index 541b1cd0e..feeed96fd 100644
--- a/dfdx/src/nn/layers/bias1d.rs
+++ b/dfdx/src/nn/layers/bias1d.rs
@@ -5,7 +5,7 @@ use crate::prelude::*;
 /// Example:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// const I: usize = 5;
 /// type Model = Bias1DConstConfig<I>;
diff --git a/dfdx/src/nn/layers/bias2d.rs b/dfdx/src/nn/layers/bias2d.rs
index 8ad2750a5..497ceb196 100644
--- a/dfdx/src/nn/layers/bias2d.rs
+++ b/dfdx/src/nn/layers/bias2d.rs
@@ -5,7 +5,7 @@ use crate::prelude::*;
 /// Example:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// const NUM_CHANS: usize = 5;
 /// type Model = Bias2DConstConfig<NUM_CHANS>;
diff --git a/dfdx/src/nn/layers/conv2d.rs b/dfdx/src/nn/layers/conv2d.rs
index 2f4648cb6..11bfb32be 100644
--- a/dfdx/src/nn/layers/conv2d.rs
+++ b/dfdx/src/nn/layers/conv2d.rs
@@ -6,7 +6,7 @@ use crate::prelude::*;
 ///
 /// Example usage:
 /// ```rust
-/// # use dfdx_nn::Conv2DConfig;
+/// # use dfdx::nn::Conv2DConfig;
 /// # use dfdx::shapes::Const;
 /// // compile time channels/kernel
 /// let m: Conv2DConfig<Const<3>, Const<5>, Const<3>> = Default::default();
diff --git a/dfdx/src/nn/layers/dropout.rs b/dfdx/src/nn/layers/dropout.rs
index 842910036..67bb9c22e 100644
--- a/dfdx/src/nn/layers/dropout.rs
+++ b/dfdx/src/nn/layers/dropout.rs
@@ -8,7 +8,7 @@ use crate::prelude::*;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// let mut dropout: DropoutOneIn<2> = Default::default();
 /// let grads = dropout.alloc_grads();
@@ -49,7 +49,7 @@ impl<const N: usize, S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> Module<Ten
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// let mut dropout = Dropout { p: 0.5 };
 /// let grads = dropout.alloc_grads();
diff --git a/dfdx/src/nn/layers/embedding.rs b/dfdx/src/nn/layers/embedding.rs
index 04f47257e..9f7767a48 100644
--- a/dfdx/src/nn/layers/embedding.rs
+++ b/dfdx/src/nn/layers/embedding.rs
@@ -17,7 +17,7 @@ use crate::prelude::*;
 ///
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = EmbeddingConstConfig<7, 2>;
 /// let mut model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/generalized_add.rs b/dfdx/src/nn/layers/generalized_add.rs
index 6fb3ef29c..6ab0b8886 100644
--- a/dfdx/src/nn/layers/generalized_add.rs
+++ b/dfdx/src/nn/layers/generalized_add.rs
@@ -10,7 +10,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = GeneralizedAdd<ReLU, Square>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/generalized_mul.rs b/dfdx/src/nn/layers/generalized_mul.rs
index 1a3c999b0..a253d5f08 100644
--- a/dfdx/src/nn/layers/generalized_mul.rs
+++ b/dfdx/src/nn/layers/generalized_mul.rs
@@ -9,7 +9,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = GeneralizedMul<ReLU, Square>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/layer_norm1d.rs b/dfdx/src/nn/layers/layer_norm1d.rs
index 2a985d8f5..1bdb1a381 100644
--- a/dfdx/src/nn/layers/layer_norm1d.rs
+++ b/dfdx/src/nn/layers/layer_norm1d.rs
@@ -13,7 +13,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = LayerNorm1DConstConfig<5>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/linear.rs b/dfdx/src/nn/layers/linear.rs
index dfc1ea3f6..74d6c711c 100644
--- a/dfdx/src/nn/layers/linear.rs
+++ b/dfdx/src/nn/layers/linear.rs
@@ -12,7 +12,7 @@ use rand_distr::Uniform;
 /// Example:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = LinearConstConfig<5, 2>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/matmul.rs b/dfdx/src/nn/layers/matmul.rs
index 0f94b0e06..7b1de6694 100644
--- a/dfdx/src/nn/layers/matmul.rs
+++ b/dfdx/src/nn/layers/matmul.rs
@@ -8,7 +8,7 @@ use rand_distr::Uniform;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = MatMulConstConfig<5, 2>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/pool_global_avg.rs b/dfdx/src/nn/layers/pool_global_avg.rs
index 28f7b44f0..599e55f91 100644
--- a/dfdx/src/nn/layers/pool_global_avg.rs
+++ b/dfdx/src/nn/layers/pool_global_avg.rs
@@ -10,7 +10,7 @@ use crate::prelude::*;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// let m: AvgPoolGlobal = Default::default();
 /// let _: Tensor<Rank1<5>, f32, _> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
diff --git a/dfdx/src/nn/layers/pool_global_max.rs b/dfdx/src/nn/layers/pool_global_max.rs
index 88d66b113..e0eea18b6 100644
--- a/dfdx/src/nn/layers/pool_global_max.rs
+++ b/dfdx/src/nn/layers/pool_global_max.rs
@@ -10,7 +10,7 @@ use crate::prelude::*;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// let m: MaxPoolGlobal = Default::default();
 /// let _: Tensor<Rank1<5>, f32, _> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
diff --git a/dfdx/src/nn/layers/pool_global_min.rs b/dfdx/src/nn/layers/pool_global_min.rs
index 4fb1df590..71f80ced8 100644
--- a/dfdx/src/nn/layers/pool_global_min.rs
+++ b/dfdx/src/nn/layers/pool_global_min.rs
@@ -10,7 +10,7 @@ use crate::prelude::*;
 /// Examples:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// let m: MinPoolGlobal = Default::default();
 /// let _: Tensor<Rank1<5>, f32, _> = m.forward(dev.zeros::<Rank3<5, 16, 8>>());
diff --git a/dfdx/src/nn/layers/reshape.rs b/dfdx/src/nn/layers/reshape.rs
index a9ee3bf00..83be2344a 100644
--- a/dfdx/src/nn/layers/reshape.rs
+++ b/dfdx/src/nn/layers/reshape.rs
@@ -5,7 +5,7 @@ use crate::prelude::*;
 /// Example usage:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// let model: Reshape<Rank2<5, 24>> = Default::default();
 /// let x: Tensor<Rank4<5, 4, 3, 2>, f32, _> = dev.sample_normal();
diff --git a/dfdx/src/nn/layers/residual_add.rs b/dfdx/src/nn/layers/residual_add.rs
index ef365fbf6..d9c707167 100644
--- a/dfdx/src/nn/layers/residual_add.rs
+++ b/dfdx/src/nn/layers/residual_add.rs
@@ -9,7 +9,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = ResidualAdd<ReLU>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/residual_mul.rs b/dfdx/src/nn/layers/residual_mul.rs
index c0598c3e4..428c1c214 100644
--- a/dfdx/src/nn/layers/residual_mul.rs
+++ b/dfdx/src/nn/layers/residual_mul.rs
@@ -8,7 +8,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = ResidualMul<ReLU>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/layers/split_into.rs b/dfdx/src/nn/layers/split_into.rs
index 45b56b356..1e66d68fe 100644
--- a/dfdx/src/nn/layers/split_into.rs
+++ b/dfdx/src/nn/layers/split_into.rs
@@ -15,7 +15,7 @@ use crate::prelude::*;
 /// # Examples
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::*;
+/// # use dfdx::*;
 /// # let dev: Cpu = Default::default();
 /// type Model = SplitInto<(LinearConstConfig<5, 3>, LinearConstConfig<5, 7>)>;
 /// let model = dev.build_module::<f32>(Model::default());
diff --git a/dfdx/src/nn/optim/adam.rs b/dfdx/src/nn/optim/adam.rs
index e0999b5bb..b94dcb7c7 100644
--- a/dfdx/src/nn/optim/adam.rs
+++ b/dfdx/src/nn/optim/adam.rs
@@ -12,11 +12,10 @@ use crate::{
 /// # Example Usage
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::{*, optim::*};
 /// # type Model = Tensor<Rank0, f32, Cpu>;
 /// # let dev: Cpu = Default::default();
 /// # let model: Model = dev.zeros();
-/// let mut opt: Adam<Model, f32, Cpu> = Adam::new(&model, AdamConfig {
+/// let mut opt: Adam<Model, f32, Cpu> = optim::Adam::new(&model, AdamConfig {
 ///     lr: 1e-2,
 ///     betas: [0.5, 0.25],
 ///     eps: 1e-6,
diff --git a/dfdx/src/nn/optim/mod.rs b/dfdx/src/nn/optim/mod.rs
index e71a57274..c4798a5ba 100644
--- a/dfdx/src/nn/optim/mod.rs
+++ b/dfdx/src/nn/optim/mod.rs
@@ -15,9 +15,8 @@
 //!
 //! ```rust
 //! # use dfdx::prelude::*;
-//! # use dfdx_nn::*;
 //! # let dev: Cpu = Default::default();
-//! use dfdx_nn::optim::*;
+//! use dfdx::nn::optim::*;
 //! type Model = LinearConstConfig<5, 2>;
 //! let mut model = dev.build_module::<f32>(Model::default());
 //! let mut grads = model.alloc_grads();
diff --git a/dfdx/src/nn/optim/rmsprop.rs b/dfdx/src/nn/optim/rmsprop.rs
index ace145c73..af7d6493b 100644
--- a/dfdx/src/nn/optim/rmsprop.rs
+++ b/dfdx/src/nn/optim/rmsprop.rs
@@ -22,11 +22,10 @@ use crate::{
 /// Constructing using new:
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::{*, optim::*};
 /// # type Model = Tensor<Rank0, f32, Cpu>;
 /// # let dev: Cpu = Default::default();
 /// # let model: Model = dev.zeros();
-/// let rmsprop: RMSprop<Model, f32, Cpu> = RMSprop::new(&model, RMSpropConfig {
+/// let rmsprop: RMSprop<Model, f32, Cpu> = optim::RMSprop::new(&model, RMSpropConfig {
 ///     lr: 1e-3,
 ///     alpha: 0.5,
 ///     eps: 1e-8,
diff --git a/dfdx/src/nn/optim/sgd.rs b/dfdx/src/nn/optim/sgd.rs
index 51d31c6cd..70671bc2a 100644
--- a/dfdx/src/nn/optim/sgd.rs
+++ b/dfdx/src/nn/optim/sgd.rs
@@ -17,11 +17,10 @@ use crate::{
 ///
 /// ```rust
 /// # use dfdx::prelude::*;
-/// # use dfdx_nn::{*, optim::*};
 /// # let dev: Cpu = Default::default();
 /// # type Model = Tensor<Rank0, f32, Cpu>;
 /// # let mut model: Model = dev.zeros();
-/// let mut opt: Sgd<Model, f32, Cpu> = Sgd::new(&model, SgdConfig {
+/// let mut opt: Sgd<Model, f32, Cpu> = optim::Sgd::new(&model, SgdConfig {
 ///     lr: 1e-3,
 ///     momentum: Some(Momentum::Classic(0.5)),
 ///     weight_decay: Some(WeightDecay::L2(0.01)),