Repository: coreylowman/dfdx (language: rust)

Issue:
Accept `impl TryInto<E>` for all scalar parameters of tensor ops
This includes:
- axpy
- clamp
- huber error
- nans to
- powf
- var/std/normalize
- scalar add
- scalar cmp
- scalar div
- scalar mul
- scalar sub


Discussion / hints:
Apparently the various builtins do not impl TryInto for other primitives (i.e. you can't TryInto<f32> for f64). Using NumCast/ToPrimitive from num_traits may work, but it may turn into breaking changes?

Fix patch:
diff --git a/src/losses.rs b/src/losses.rs
index 7f9923014..e09ef67fc 100644
--- a/src/losses.rs
+++ b/src/losses.rs
@@ -47,7 +47,7 @@ pub fn mae_loss<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>>(
 pub fn huber_loss<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>>(
     pred: Tensor<S, E, D, T>,
     targ: Tensor<S, E, D>,
-    delta: E,
+    delta: impl Into<E>,
 ) -> Tensor<Rank0, E, D, T> {
     pred.huber_error(targ, delta).mean()
 }
@@ -62,8 +62,9 @@ pub fn huber_loss<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>>(
 pub fn smooth_l1_loss<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>>(
     pred: Tensor<S, E, D, T>,
     targ: Tensor<S, E, D>,
-    delta: E,
+    delta: impl Into<E>,
 ) -> Tensor<Rank0, E, D, T> {
+    let delta = delta.into();
     huber_loss(pred, targ, delta) / delta
 }
 
diff --git a/src/tensor_ops/add/mod.rs b/src/tensor_ops/add/mod.rs
index 86a10044f..402cf8952 100644
--- a/src/tensor_ops/add/mod.rs
+++ b/src/tensor_ops/add/mod.rs
@@ -3,7 +3,7 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
-use super::{ops::*, Device};
+use super::ops::*;
 use crate::{
     shapes::*,
     tensor::{DeviceStorage, HasErr, Merge, Tape, Tensor},
@@ -38,10 +38,13 @@ pub struct ScalarAddKernelOp<E> {
 /// let r = a + 1.0;
 /// assert_eq!(r.array(), [[2.0, 3.0, 4.0], [0.0, -1.0, -2.0]]);
 /// ```
-pub fn add<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<R>, R: Default>(
+pub fn add<S: Shape, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Default>(
     lhs: Tensor<S, E, D, T>,
     rhs: Tensor<S, E, D, R>,
-) -> Tensor<S, E, D, T> {
+) -> Tensor<S, E, D, T>
+where
+    D: BinaryKernel<BinaryAddKernelOp, E>,
+{
     lhs + rhs
 }
 
diff --git a/src/tensor_ops/axpy/mod.rs b/src/tensor_ops/axpy/mod.rs
index 73c72df38..8299eb2ae 100644
--- a/src/tensor_ops/axpy/mod.rs
+++ b/src/tensor_ops/axpy/mod.rs
@@ -12,9 +12,9 @@ mod cuda_kernel;
 /// See [Tensor::axpy] for in place version.
 pub fn axpy<S: Shape, E: Unit, D>(
     a: &Tensor<S, E, D>,
-    alpha: E,
+    alpha: impl Into<E>,
     b: &Tensor<S, E, D>,
-    beta: E,
+    beta: impl Into<E>,
 ) -> Tensor<S, E, D>
 where
     D: AxpyKernel<E>,
@@ -26,19 +26,24 @@ where
 
 impl<S: Shape, E: Unit, D: AxpyKernel<E>> Tensor<S, E, D> {
     /// Updates self with elementwise function `self = self * alpha + b * beta`.
-    pub fn axpy<T>(&mut self, alpha: E, b: &Tensor<S, E, D, T>, beta: E) {
+    pub fn axpy<T>(&mut self, alpha: impl Into<E>, b: &Tensor<S, E, D, T>, beta: impl Into<E>) {
         self.try_axpy(alpha, b, beta).unwrap()
     }
 
     /// Updates self with elementwise function `self = self * alpha + b * beta`.
-    pub fn try_axpy<T>(&mut self, alpha: E, b: &Tensor<S, E, D, T>, beta: E) -> Result<(), D::Err> {
+    pub fn try_axpy<T>(
+        &mut self,
+        alpha: impl Into<E>,
+        b: &Tensor<S, E, D, T>,
+        beta: impl Into<E>,
+    ) -> Result<(), D::Err> {
         assert_eq!(self.shape, b.shape);
         assert_eq!(self.strides, b.strides, "Strides must be equal for axpy");
         self.device.clone().forward(
             std::sync::Arc::make_mut(&mut self.data),
-            alpha,
+            alpha.into(),
             b.data.as_ref(),
-            beta,
+            beta.into(),
         )
     }
 }
diff --git a/src/tensor_ops/clamp/mod.rs b/src/tensor_ops/clamp/mod.rs
index 2d75b51ef..ae7901fc4 100644
--- a/src/tensor_ops/clamp/mod.rs
+++ b/src/tensor_ops/clamp/mod.rs
@@ -25,20 +25,26 @@ pub struct ClampKernelOp<E> {
 /// ```
 pub fn clamp<S: Shape, E: Dtype, D: UnaryKernel<ClampKernelOp<E>, E>, T: Tape<E, D>>(
     t: Tensor<S, E, D, T>,
-    min: E,
-    max: E,
+    min: impl Into<E>,
+    max: impl Into<E>,
 ) -> Tensor<S, E, D, T> {
     t.clamp(min, max)
 }
 
 impl<S: Shape, E: Dtype, D: UnaryKernel<ClampKernelOp<E>, E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     /// See [clamp]
-    pub fn clamp(self, min: E, max: E) -> Self {
+    pub fn clamp(self, min: impl Into<E>, max: impl Into<E>) -> Self {
         self.try_clamp(min, max).unwrap()
     }
     /// See [clamp]
-    pub fn try_clamp(self, min: E, max: E) -> Result<Self, D::Err> {
-        try_unary_op(ClampKernelOp { min, max }, self)
+    pub fn try_clamp(self, min: impl Into<E>, max: impl Into<E>) -> Result<Self, D::Err> {
+        try_unary_op(
+            ClampKernelOp {
+                min: min.into(),
+                max: max.into(),
+            },
+            self,
+        )
     }
 }
 
diff --git a/src/tensor_ops/cmp/mod.rs b/src/tensor_ops/cmp/mod.rs
index 7052cebd1..cdb7a0adc 100644
--- a/src/tensor_ops/cmp/mod.rs
+++ b/src/tensor_ops/cmp/mod.rs
@@ -216,13 +216,13 @@ macro_rules! impl_cmp_kernel_op {
             #[doc = $doc]
             pub fn $try_scalar_op(
                 &self,
-                scalar: E,
+                scalar: impl Into<E>,
             ) -> Result<Tensor<S, bool, D, NoneTape>, D::Err> {
-                try_scalar_cmp_op(self, scalar)
+                try_scalar_cmp_op(self, scalar.into())
             }
 
             #[doc = $doc]
-            pub fn $scalar_op(&self, scalar: E) -> Tensor<S, bool, D, NoneTape> {
+            pub fn $scalar_op(&self, scalar: impl Into<E>) -> Tensor<S, bool, D, NoneTape> {
                 self.$try_scalar_op(scalar).unwrap()
             }
         }
diff --git a/src/tensor_ops/div/mod.rs b/src/tensor_ops/div/mod.rs
index 4ce5bc74d..f7abb2ead 100644
--- a/src/tensor_ops/div/mod.rs
+++ b/src/tensor_ops/div/mod.rs
@@ -3,7 +3,7 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
-use super::{ops::*, Device};
+use super::ops::*;
 use crate::{shapes::*, tensor::*};
 
 #[repr(C)]
@@ -36,10 +36,13 @@ pub struct BinaryDivKernelOp;
 /// let r = a / 2.0;
 /// assert_eq!(r.array(), [[0.5, 1.0, 1.5], [-0.5, -1.0, -1.5]]);
 /// ```
-pub fn div<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<R>, R: Default>(
+pub fn div<S: Shape, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Default>(
     lhs: Tensor<S, E, D, T>,
     rhs: Tensor<S, E, D, R>,
-) -> Tensor<S, E, D, T> {
+) -> Tensor<S, E, D, T>
+where
+    D: BinaryKernel<BinaryDivKernelOp, E>,
+{
     lhs / rhs
 }
 
@@ -48,9 +51,10 @@ pub trait TryDiv<Rhs = Self>: HasErr {
     fn try_div(self, rhs: Rhs) -> Result<Self, Self::Err>;
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, LhsTape: Tape<E, D>, R> TryDiv<Tensor<S, E, D, R>>
+impl<S: Shape, E: Dtype, D, LhsTape: Tape<E, D>, R> TryDiv<Tensor<S, E, D, R>>
     for Tensor<S, E, D, LhsTape>
 where
+    D: BinaryKernel<BinaryDivKernelOp, E>,
     LhsTape: Merge<R>,
 {
     /// See [div]
@@ -59,14 +63,16 @@ where
     }
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> TryDiv<E> for Tensor<S, E, D, T> {
+impl<S: Shape, E: Dtype, D: UnaryKernel<ScalarDivKernelOp<E>, E>, T: Tape<E, D>> TryDiv<E>
+    for Tensor<S, E, D, T>
+{
     /// See [div]
     fn try_div(self, rhs: E) -> Result<Self, Self::Err> {
         try_unary_op(ScalarDivKernelOp { scalar: rhs }, self)
     }
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, LhsTape: Tape<E, D>, Rhs> std::ops::Div<Rhs>
+impl<S: Shape, E: Dtype, D: DeviceStorage, LhsTape: Tape<E, D>, Rhs> std::ops::Div<Rhs>
     for Tensor<S, E, D, LhsTape>
 where
     Self: TryDiv<Rhs>,
diff --git a/src/tensor_ops/dropout/mod.rs b/src/tensor_ops/dropout/mod.rs
index a11352c3d..890ac3282 100644
--- a/src/tensor_ops/dropout/mod.rs
+++ b/src/tensor_ops/dropout/mod.rs
@@ -51,19 +51,20 @@ pub trait DropoutKernel<E: Dtype>: DeviceStorage {
 /// random numbers, so the masking is the same for both.
 pub fn dropout<S: Shape, E: Dtype, D: DropoutKernel<E>, T: Tape<E, D>>(
     t: Tensor<S, E, D, T>,
-    prob: E,
+    prob: impl Into<E>,
 ) -> Tensor<S, E, D, T> {
     t.dropout(prob)
 }
 
 impl<S: Shape, E: Dtype, D: DropoutKernel<E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     /// See [dropout]
-    pub fn dropout(self, prob: E) -> Self {
+    pub fn dropout(self, prob: impl Into<E>) -> Self {
         self.try_dropout(prob).unwrap()
     }
     /// See [dropout]
-    pub fn try_dropout(self, prob: E) -> Result<Self, D::Err> {
+    pub fn try_dropout(self, prob: impl Into<E>) -> Result<Self, D::Err> {
         let seed = self.device.random_u64();
+        let prob = prob.into();
         let op = DropoutKernelOp { seed, prob };
         let (inp, mut tape) = self.split_tape();
         let out = inp.device.forward(op, &inp)?;
diff --git a/src/tensor_ops/huber_error/mod.rs b/src/tensor_ops/huber_error/mod.rs
index c64d77950..c45a60f69 100644
--- a/src/tensor_ops/huber_error/mod.rs
+++ b/src/tensor_ops/huber_error/mod.rs
@@ -31,14 +31,14 @@ pub struct HuberErrorKernelOp<E> {
 pub fn huber_error<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<R>, R: Tape<E, D>>(
     lhs: Tensor<S, E, D, T>,
     rhs: Tensor<S, E, D, R>,
-    delta: E,
+    delta: impl Into<E>,
 ) -> Tensor<S, E, D, T> {
     lhs.huber_error(rhs, delta)
 }
 
 impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     /// See [huber_error]
-    pub fn huber_error<R: Tape<E, D>>(self, rhs: Tensor<S, E, D, R>, delta: E) -> Self
+    pub fn huber_error<R: Tape<E, D>>(self, rhs: Tensor<S, E, D, R>, delta: impl Into<E>) -> Self
     where
         T: Merge<R>,
     {
@@ -49,11 +49,12 @@ impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     pub fn try_huber_error<R: Tape<E, D>>(
         self,
         rhs: Tensor<S, E, D, R>,
-        delta: E,
+        delta: impl Into<E>,
     ) -> Result<Self, D::Err>
     where
         T: Merge<R>,
     {
+        let delta = delta.into();
         try_binary_op(HuberErrorKernelOp { delta }, self, rhs)
     }
 }
diff --git a/src/tensor_ops/mul/mod.rs b/src/tensor_ops/mul/mod.rs
index 33780b942..ff45b2c47 100644
--- a/src/tensor_ops/mul/mod.rs
+++ b/src/tensor_ops/mul/mod.rs
@@ -3,7 +3,7 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
-use super::{ops::*, Device};
+use super::ops::*;
 use crate::{shapes::*, tensor::*};
 
 #[repr(C)]
@@ -35,10 +35,13 @@ pub struct ScalarMulKernelOp<E> {
 /// let r = a * 2.0;
 /// assert_eq!(r.array(), [[2.0, 4.0, 6.0], [-2.0, -4.0, -6.0]]);
 /// ```
-pub fn mul<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<R>, R: Default>(
+pub fn mul<S: Shape, E: Dtype, D, T: Tape<E, D> + Merge<R>, R: Default>(
     lhs: Tensor<S, E, D, T>,
     rhs: Tensor<S, E, D, R>,
-) -> Tensor<S, E, D, T> {
+) -> Tensor<S, E, D, T>
+where
+    D: BinaryKernel<BinaryMulKernelOp, E>,
+{
     lhs * rhs
 }
 
@@ -47,8 +50,8 @@ pub trait TryMul<Rhs = Self>: HasErr {
     fn try_mul(self, rhs: Rhs) -> Result<Self, Self::Err>;
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, LhsTape: Tape<E, D>, R> TryMul<Tensor<S, E, D, R>>
-    for Tensor<S, E, D, LhsTape>
+impl<S: Shape, E: Dtype, D: BinaryKernel<BinaryMulKernelOp, E>, LhsTape: Tape<E, D>, R>
+    TryMul<Tensor<S, E, D, R>> for Tensor<S, E, D, LhsTape>
 where
     LhsTape: Merge<R>,
 {
@@ -57,13 +60,15 @@ where
     }
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> TryMul<E> for Tensor<S, E, D, T> {
+impl<S: Shape, E: Dtype, D: UnaryKernel<ScalarMulKernelOp<E>, E>, T: Tape<E, D>> TryMul<E>
+    for Tensor<S, E, D, T>
+{
     fn try_mul(self, rhs: E) -> Result<Self, Self::Err> {
         try_unary_op(ScalarMulKernelOp { scalar: rhs }, self)
     }
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, LhsTape: Tape<E, D>, Rhs> std::ops::Mul<Rhs>
+impl<S: Shape, E: Dtype, D: DeviceStorage, LhsTape: Tape<E, D>, Rhs> std::ops::Mul<Rhs>
     for Tensor<S, E, D, LhsTape>
 where
     Self: TryMul<Rhs>,
diff --git a/src/tensor_ops/nans_to/mod.rs b/src/tensor_ops/nans_to/mod.rs
index 7167b6189..fa98308da 100644
--- a/src/tensor_ops/nans_to/mod.rs
+++ b/src/tensor_ops/nans_to/mod.rs
@@ -24,18 +24,19 @@ pub struct NansToKernelOp<E>(E);
 /// ```
 pub fn nans_to<S: Shape, E: Dtype, D: UnaryKernel<NansToKernelOp<E>, E>, T: Tape<E, D>>(
     t: Tensor<S, E, D, T>,
-    value: E,
+    value: impl Into<E>,
 ) -> Tensor<S, E, D, T> {
     t.nans_to(value)
 }
 
 impl<S: Shape, E: Dtype, D: UnaryKernel<NansToKernelOp<E>, E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     /// See [nans_to]
-    pub fn nans_to(self, value: E) -> Self {
+    pub fn nans_to(self, value: impl Into<E>) -> Self {
         self.try_nans_to(value).unwrap()
     }
     /// See [nans_to]
-    pub fn try_nans_to(self, value: E) -> Result<Self, D::Err> {
+    pub fn try_nans_to(self, value: impl Into<E>) -> Result<Self, D::Err> {
+        let value = value.into();
         try_unary_op(NansToKernelOp(value), self)
     }
 }
diff --git a/src/tensor_ops/normalize.rs b/src/tensor_ops/normalize.rs
index e3657bab5..2da3db0ff 100644
--- a/src/tensor_ops/normalize.rs
+++ b/src/tensor_ops/normalize.rs
@@ -17,14 +17,14 @@ use super::{BroadcastTo, Device, MeanTo, TryAdd, TryDiv, TrySub};
 /// ```
 pub fn normalize<Ax: Axes, S: Shape + ReduceShape<Ax>, E: Dtype, D: Device<E>, T: Tape<E, D>>(
     t: Tensor<S, E, D, T>,
-    epsilon: E,
+    epsilon: impl Into<E>,
 ) -> Tensor<S, E, D, T> {
     t.normalize::<Ax>(epsilon)
 }
 
 impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     /// See [normalize]
-    pub fn normalize<Ax: Axes>(self, epsilon: E) -> Self
+    pub fn normalize<Ax: Axes>(self, epsilon: impl Into<E>) -> Self
     where
         S: ReduceShape<Ax>,
     {
@@ -32,7 +32,10 @@ impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     }
 
     /// See [normalize]
-    pub fn try_normalize<Ax: Axes>(self, epsilon: E) -> Result<Self, <Self as HasErr>::Err>
+    pub fn try_normalize<Ax: Axes>(
+        self,
+        epsilon: impl Into<E>,
+    ) -> Result<Self, <Self as HasErr>::Err>
     where
         S: ReduceShape<Ax>,
     {
@@ -43,7 +46,7 @@ impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> Tensor<S, E, D, T> {
             .retaped::<T>()
             .try_square()?
             .try_mean::<_, Ax>()?
-            .try_add(epsilon)?
+            .try_add(epsilon.into())?
             .try_sqrt()?;
         centered.try_div(std.try_broadcast_like(&shape)?)
     }
diff --git a/src/tensor_ops/pow/mod.rs b/src/tensor_ops/pow/mod.rs
index 7d8db4a70..6cdc24580 100644
--- a/src/tensor_ops/pow/mod.rs
+++ b/src/tensor_ops/pow/mod.rs
@@ -23,18 +23,19 @@ pub struct PowfKernelOp<E>(E);
 /// ```
 pub fn powf<S: Shape, E: Dtype, D: UnaryKernel<PowfKernelOp<E>, E>, T: Tape<E, D>>(
     t: Tensor<S, E, D, T>,
-    exponent: E,
+    exponent: impl Into<E>,
 ) -> Tensor<S, E, D, T> {
     t.powf(exponent)
 }
 
 impl<S: Shape, E: Dtype, D: UnaryKernel<PowfKernelOp<E>, E>, T: Tape<E, D>> Tensor<S, E, D, T> {
     /// See [powf]
-    pub fn powf(self, exponent: E) -> Self {
+    pub fn powf(self, exponent: impl Into<E>) -> Self {
         self.try_powf(exponent).unwrap()
     }
     /// See [powf]
-    pub fn try_powf(self, exponent: E) -> Result<Self, D::Err> {
+    pub fn try_powf(self, exponent: impl Into<E>) -> Result<Self, D::Err> {
+        let exponent = exponent.into();
         try_unary_op(PowfKernelOp(exponent), self)
     }
 }
diff --git a/src/tensor_ops/stddev_to.rs b/src/tensor_ops/stddev_to.rs
index 34962a951..5b7c03da1 100644
--- a/src/tensor_ops/stddev_to.rs
+++ b/src/tensor_ops/stddev_to.rs
@@ -15,7 +15,7 @@ pub trait StddevTo<E: Dtype>: HasErr + HasShape {
     /// let r = t.stddev::<Rank1<2>, _>(0.0); // or `stddev::<_, Axis<1>>(0.0)`
     /// assert_eq!(r.array(), [0.6666667_f32.sqrt(), 6.0_f32.sqrt()]);
     /// ```
-    fn stddev<Dst: Shape, Ax: Axes>(self, epsilon: E) -> Self::WithShape<Dst>
+    fn stddev<Dst: Shape, Ax: Axes>(self, epsilon: impl Into<E>) -> Self::WithShape<Dst>
     where
         Self::Shape: HasAxes<Ax> + ReduceShapeTo<Dst, Ax>,
     {
@@ -24,18 +24,21 @@ pub trait StddevTo<E: Dtype>: HasErr + HasShape {
     /// Fallible version of [StddevTo::stddev]
     fn try_stddev<Dst: Shape, Ax: Axes>(
         self,
-        epsilon: E,
+        epsilon: impl Into<E>,
     ) -> Result<Self::WithShape<Dst>, Self::Err>
     where
         Self::Shape: HasAxes<Ax> + ReduceShapeTo<Dst, Ax>;
 }
 
 impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> StddevTo<E> for Tensor<S, E, D, T> {
-    fn try_stddev<Dst: Shape, Ax: Axes>(self, epsilon: E) -> Result<Self::WithShape<Dst>, Self::Err>
+    fn try_stddev<Dst: Shape, Ax: Axes>(
+        self,
+        epsilon: impl Into<E>,
+    ) -> Result<Self::WithShape<Dst>, Self::Err>
     where
         Self::Shape: HasAxes<Ax> + ReduceShapeTo<Dst, Ax>,
     {
-        self.try_var()?.try_add(epsilon)?.try_sqrt()
+        self.try_var()?.try_add(epsilon.into())?.try_sqrt()
     }
 }
 
diff --git a/src/tensor_ops/sub/mod.rs b/src/tensor_ops/sub/mod.rs
index 573dac7bb..a54c8bd70 100644
--- a/src/tensor_ops/sub/mod.rs
+++ b/src/tensor_ops/sub/mod.rs
@@ -3,7 +3,7 @@ mod cpu_kernel;
 #[cfg(feature = "cuda")]
 mod cuda_kernel;
 
-use super::{ops::*, Device};
+use super::ops::*;
 use crate::{shapes::*, tensor::*};
 
 #[repr(C)]
@@ -36,10 +36,13 @@ pub struct ScalarSubKernelOp<E> {
 /// let r = a - 1.0;
 /// assert_eq!(r.array(), [[0.0, 1.0, 2.0], [-2.0, -3.0, -4.0]]);
 /// ```
-pub fn sub<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D> + Merge<R>, R>(
+pub fn sub<S: Shape, E: Dtype, D, T: Tape<E, D> + Merge<R>, R>(
     lhs: Tensor<S, E, D, T>,
     rhs: Tensor<S, E, D, R>,
-) -> Tensor<S, E, D, T> {
+) -> Tensor<S, E, D, T>
+where
+    D: BinaryKernel<BinarySubKernelOp, E>,
+{
     lhs - rhs
 }
 
@@ -48,8 +51,8 @@ pub trait TrySub<Rhs = Self>: HasErr {
     fn try_sub(self, rhs: Rhs) -> Result<Self, Self::Err>;
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, LTape: Tape<E, D>, R> TrySub<Tensor<S, E, D, R>>
-    for Tensor<S, E, D, LTape>
+impl<S: Shape, E: Dtype, D: BinaryKernel<BinarySubKernelOp, E>, LTape: Tape<E, D>, R>
+    TrySub<Tensor<S, E, D, R>> for Tensor<S, E, D, LTape>
 where
     LTape: Merge<R>,
 {
@@ -58,13 +61,15 @@ where
     }
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, T: Tape<E, D>> TrySub<E> for Tensor<S, E, D, T> {
+impl<S: Shape, E: Dtype, D: UnaryKernel<ScalarSubKernelOp<E>, E>, T: Tape<E, D>> TrySub<E>
+    for Tensor<S, E, D, T>
+{
     fn try_sub(self, rhs: E) -> Result<Self, Self::Err> {
         try_unary_op(ScalarSubKernelOp { scalar: rhs }, self)
     }
 }
 
-impl<S: Shape, E: Dtype, D: Device<E>, LTape: Tape<E, D>, Rhs> std::ops::Sub<Rhs>
+impl<S: Shape, E: Dtype, D: DeviceStorage, LTape: Tape<E, D>, Rhs> std::ops::Sub<Rhs>
     for Tensor<S, E, D, LTape>
 where
     Self: TrySub<Rhs>,


Test patch: