fix: Fix wrong condition in smooth l1 loss (#213)

* fix wrong condition in smooth l1 loss

* clippy & fmt

* add unit test for smooth l1 loss with negative diff

* clippy and bump up patch version

* update unit test

Author: minghuaw

Cargo.toml

@@ -1,7 +1,7 @@
 [workspace.package]
 # All but mlx-sys should follow the same version. mlx-sys should follow 
 # the version of mlx-c.
-version = "0.21.1"
+version = "0.21.2"
 edition = "2021"
 authors = [
     "Minghua Wu <michael.wu1107@gmail.com>",
@@ -29,9 +29,9 @@ resolver = "2"
 [workspace.dependencies]
 # workspace local dependencies
 mlx-sys = { version = "=0.1.0", path = "mlx-sys" }
-mlx-macros = { version = "0.21.0", path = "mlx-macros" }
-mlx-internal-macros = { version = "0.21.0", path = "mlx-internal-macros" }
-mlx-rs = { version = "0.21.0", path = "mlx-rs" }
+mlx-macros = { version = "0.21", path = "mlx-macros" }
+mlx-internal-macros = { version = "0.21", path = "mlx-internal-macros" }
+mlx-rs = { version = "0.21.2", path = "mlx-rs" }
 
 # external dependencies
 thiserror = "1"

mlx-rs/src/losses.rs

@@ -523,13 +523,12 @@ impl SmoothL1Loss {
         let reduction = self.reduction;
 
         check_shape(predictions, targets, "predictions", "targets")?;
-        let diff = predictions.subtract(targets)?;
+        let diff = predictions.subtract(targets)?.abs()?;
+        let beta = array!(beta);
         let loss = r#where(
-            &diff.lt(array!(beta))?,
-            array!(0.5)
-                .multiply(square(&diff)?)?
-                .divide(&array!(beta))?,
-            abs(&diff)?.subtract(array!(0.5).multiply(array!(beta))?)?,
+            &diff.lt(&beta)?,
+            array!(0.5).multiply(square(&diff)?)?.divide(&beta)?,
+            diff.subtract(array!(0.5).multiply(beta)?)?,
         )?;
         reduction.reduce(loss)
     }
@@ -1189,6 +1188,18 @@ mod tests {
         assert_array_eq!(loss_mean, expected_mean);
     }
 
+    #[test]
+    fn test_smooth_l1_loss_negative_diff() {
+        let a = array!([1.5, 6.0, 0.5, 2.5]);
+        let b = array!([1.0, 2.0, 0.5, 3.5]);
+
+        let loss = SmoothL1Loss::new();
+
+        let ab = loss.apply(&a, &b).unwrap();
+        let ba = loss.apply(&b, &a).unwrap();
+        assert_array_eq!(ab, ba);
+    }
+
     #[test]
     fn test_nll_loss() {
         let logits = array!([[0.0, f32::NEG_INFINITY], [f32::NEG_INFINITY, 0.0]]);

mlx-rs/src/transforms/mod.rs

@@ -72,21 +72,22 @@ fn jvp_inner(
 
 /// Compute the Jacobian-vector product.
 ///
-/// This computes the product of the Jacobian of a function `f` evaluated at `primals` with the
-/// `tangents`.
+/// This computes the product of the Jacobian of a function `f` evaluated at
+/// `primals` with the `tangents`.
 ///
 /// # Params:
 ///
-/// - `f`: function which takes an array of `Array` and returns an array of `Array`
+/// - `f`: function which takes an array of `Array` and returns an array of
+///   `Array`
 /// - `primals`: array of `Array` at which to evaluate the Jacobian
-/// - `tangents`: array of `Array` which are the "vector" in the Jacobian-vector product.  The
-///     `tangents` should be the same in number, shape and type as the inputs of `f`, e.g. the
-///     `primals`
+/// - `tangents`: array of `Array` which are the "vector" in the Jacobian-vector
+///   product.  The `tangents` should be the same in number, shape and type as
+///   the inputs of `f`, e.g. the `primals`
 ///
 /// # Returns:
 ///
-/// Array of the Jacobian-vector products which is the same in number, shape and type of
-/// the outputs of `f`
+/// Array of the Jacobian-vector products which is the same in number, shape and
+/// type of the outputs of `f`
 pub fn jvp<'a, F>(f: F, primals: &[Array], tangents: &[Array]) -> Result<(Vec<Array>, Vec<Array>)>
 where
     F: FnMut(&[Array]) -> Vec<Array> + 'a,