mlx_rs/
losses.rs

1//! Loss functions
2
3use crate::{
4    array,
5    error::{CrossEntropyBuildError, Exception},
6    ops::{
7        abs, clip, exp, indexing::take_along_axis, log, logaddexp, logsumexp_axes, maximum,
8        minimum, multiply, power, select, sqrt, square, sum_axes, sum_axis,
9    },
10    Array,
11};
12use mlx_internal_macros::{generate_builder, Buildable};
13
14#[inline]
15fn check_shape(
16    left: &Array,
17    right: &Array,
18    left_ident: &str,
19    right_ident: &str,
20) -> Result<(), Exception> {
21    if left.shape() != right.shape() {
22        return Err(Exception::custom(format!(
23            "The shape of the {} ({:?}) does not match the shape of the {} ({:?})",
24            left_ident,
25            left.shape(),
26            right_ident,
27            right.shape()
28        )));
29    }
30    Ok(())
31}
32
33/// Different types of loss reductions
34#[derive(Debug, Clone, Copy)]
35pub enum LossReduction {
36    /// No reduction is applied.
37    None,
38    /// The sum of the output will be computed.
39    Sum,
40    /// The mean of the output will be computed.
41    Mean,
42}
43
44impl LossReduction {
45    /// Reduces the loss according to the reduction type.
46    pub fn reduce(&self, loss: Array) -> Result<Array, Exception> {
47        match self {
48            LossReduction::None => Ok(loss),
49            LossReduction::Sum => Ok(loss.sum(None)?),
50            LossReduction::Mean => Ok(loss.mean(None)?),
51        }
52    }
53}
54
55/// Helper type alias for CrossEntropyBuilder weights.
56pub type CrossEntropyBuilderWeights<'a> = &'a Array;
57
58generate_builder! {
59    /// Cross entropy loss function.
60    #[derive(Debug, Clone, Buildable)]
61    #[buildable(root = crate)]
62    #[builder(
63        root = crate,
64        build_with = build_cross_entropy,
65        err = CrossEntropyBuildError
66    )]
67    pub struct CrossEntropy<'a> {
68        /// Weights for each target
69        #[builder(optional, default = CrossEntropy::DEFAULT_WEIGHTS)]
70        pub weights: Option<&'a Array>,
71
72        /// The axis over which to compute softmax. Default to [`CrossEntropy::DEFAULT_AXIS`]
73        #[builder(optional, default = CrossEntropy::DEFAULT_AXIS)]
74        pub axis: i32,
75
76        /// The label smoothing factor, range [0, 1). Default to
77        /// [`CrossEntropy::DEFAULT_LABEL_SMOOTHING`]
78        #[builder(optional, default = CrossEntropy::DEFAULT_LABEL_SMOOTHING)]
79        pub label_smoothing: f32,
80
81        /// Reduction type. Default to [`CrossEntropy::DEFAULT_REDUCTION`]
82        #[builder(optional, default = CrossEntropy::DEFAULT_REDUCTION)]
83        pub reduction: LossReduction,
84    }
85}
86
87fn build_cross_entropy(
88    builder: CrossEntropyBuilder,
89) -> Result<CrossEntropy, CrossEntropyBuildError> {
90    let axis = builder.axis;
91    let label_smoothing = builder.label_smoothing;
92    let reduction = builder.reduction;
93
94    if !(0.0..1.0).contains(&label_smoothing) {
95        return Err(CrossEntropyBuildError::InvalidLabelSmoothingFactor);
96    }
97
98    Ok(CrossEntropy {
99        weights: builder.weights,
100        axis,
101        label_smoothing,
102        reduction,
103    })
104}
105
106impl<'a> CrossEntropy<'a> {
107    /// Default value for the `axis` parameter.
108    pub const DEFAULT_AXIS: i32 = -1;
109
110    /// Default value for the `label_smoothing` parameter.
111    pub const DEFAULT_LABEL_SMOOTHING: f32 = 0.0;
112
113    /// Default value for the `reduction` parameter.
114    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
115
116    /// Default value for the `weights` parameter.
117    pub const DEFAULT_WEIGHTS: Option<&'a Array> = None;
118
119    /// Apply the cross entropy loss function on the given logits and targets.
120    ///
121    /// # Params
122    ///
123    /// - `logits`: unnormalized predicted logits
124    /// - `targets`: target values, as class indices
125    pub fn apply(
126        &self,
127        logits: impl AsRef<Array>,
128        targets: impl AsRef<Array>,
129    ) -> Result<Array, Exception> {
130        let logits = logits.as_ref();
131        let targets = targets.as_ref();
132
133        let target_as_probs = targets.ndim() == logits.ndim();
134
135        let score = if target_as_probs {
136            sum_axes(&logits.multiply(targets)?, &[self.axis], None)?
137        } else {
138            take_along_axis(logits, &targets.expand_dims_axes(&[-1])?, self.axis)?
139                .squeeze_axes(&[-1])?
140        };
141        let log_sum_exp_logits = logsumexp_axes(logits, &[self.axis], None)?;
142
143        let mut loss = if self.label_smoothing > 0.0 {
144            // adjust the true class score with label smoothing
145            let adjusted_score = multiply(array!(1.0 - self.label_smoothing), score)?;
146
147            // calculate the mean logit across the classes for smoothed loss
148            let mean_logits = logits.mean_axis(self.axis, None)?;
149            let smoothed_loss = -multiply(mean_logits, array!(self.label_smoothing))?;
150
151            // combine the adjusted score and smoothed loss with the logsumexp logits
152            log_sum_exp_logits
153                .subtract(adjusted_score)?
154                .add(smoothed_loss)?
155        } else {
156            log_sum_exp_logits.subtract(score)?
157        };
158
159        if let Some(weights) = self.weights {
160            check_shape(weights, &loss, "weights", "loss")?;
161            loss = multiply(loss, weights)?;
162        }
163
164        self.reduction.reduce(loss)
165    }
166}
167
168generate_builder! {
169    /// Binary cross entropy loss.
170    ///
171    /// By default, this function takes the pre-sigmoid logits, which results in a faster
172    /// and more precise loss. For improved numerical stability when `inputs_are_logits` is true,
173    /// the loss calculation clips the input probabilities (in log-space) to a minimum value
174    /// of `-100`.
175    #[derive(Debug, Clone, Buildable)]
176    #[buildable(root = crate)]
177    #[builder(root = crate)]
178    pub struct BinaryCrossEntropy<'a> {
179        /// Optional weights for each target
180        #[builder(optional, default = BinaryCrossEntropy::DEFAULT_WEIGHTS)]
181        pub weights: Option<&'a Array>,
182
183        /// Whether the inputs are logits. Default to
184        /// [`BinaryCrossEntropy::DEFAULT_INPUTS_ARE_LOGITS`]
185        #[builder(optional, default = BinaryCrossEntropy::DEFAULT_INPUTS_ARE_LOGITS)]
186        pub inputs_are_logits: bool,
187
188        /// Reduction type. Default to [`BinaryCrossEntropy::DEFAULT_REDUCTION`]
189        #[builder(optional, default = BinaryCrossEntropy::DEFAULT_REDUCTION)]
190        pub reduction: LossReduction,
191    }
192}
193
194impl<'a> BinaryCrossEntropy<'a> {
195    /// Default value for the `weights` parameter.
196    pub const DEFAULT_WEIGHTS: Option<&'a Array> = None;
197
198    /// Default value for the `with_logits` parameter.
199    pub const DEFAULT_INPUTS_ARE_LOGITS: bool = true;
200
201    /// Default value for the `reduction` parameter.
202    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
203
204    /// Apply the binary cross entropy loss function on the given logits and targets.
205    ///
206    /// # Params
207    ///
208    /// - `logits`: unnormalized predicted logits
209    /// - `targets`: binary target values in {0, 1}
210    pub fn apply(
211        &self,
212        logits: impl AsRef<Array>,
213        targets: impl AsRef<Array>,
214    ) -> Result<Array, Exception> {
215        let logits = logits.as_ref();
216        let targets = targets.as_ref();
217        let weights = self.weights;
218        let inputs_are_logits = self.inputs_are_logits;
219        let reduction = self.reduction;
220
221        let mut loss = if inputs_are_logits {
222            logaddexp(array!(0.0), logits)?.subtract(targets.multiply(logits)?)?
223        } else {
224            let log_inputs_clip = clip(log(logits)?, (-100.0, ()))?;
225            let log_inputs_inverse_clip = clip(log(&array!(1.0).subtract(logits)?)?, (-100.0, ()))?;
226            -(targets.multiply(log_inputs_clip)?.add(
227                array!(1.0)
228                    .subtract(targets)?
229                    .multiply(log_inputs_inverse_clip)?,
230            )?)
231        };
232
233        if let Some(weights) = weights {
234            check_shape(weights, &loss, "weights", "loss")?;
235            loss = multiply(loss, weights)?;
236        }
237
238        reduction.reduce(loss)
239    }
240}
241
242generate_builder! {
243    /// Computes the L1 loss
244    #[derive(Debug, Clone, Buildable)]
245    #[buildable(root = crate)]
246    #[builder(root = crate)]
247    pub struct L1Loss {
248        /// Reduction type. Default to [`L1loss::DEFAULT_REDUCTION`]
249        #[builder(optional, default = L1Loss::DEFAULT_REDUCTION)]
250        pub reduction: LossReduction,
251    }
252}
253
254impl L1Loss {
255    /// Default value for the `reduction` parameter.
256    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::Mean;
257
258    /// Compute the L1 loss.
259    ///
260    /// # Params
261    ///
262    /// - `predictions`: predicted values
263    /// - `targets`: target values
264    pub fn apply(
265        &self,
266        predictions: impl AsRef<Array>,
267        targets: impl AsRef<Array>,
268    ) -> Result<Array, Exception> {
269        let predictions = predictions.as_ref();
270        let targets = targets.as_ref();
271        let reduction = self.reduction;
272
273        check_shape(predictions, targets, "predictions", "targets")?;
274        let loss = predictions.subtract(targets)?.abs()?;
275        reduction.reduce(loss)
276    }
277}
278
279generate_builder! {
280    /// Computes the mean squared error loss.
281    #[derive(Debug, Clone, Buildable)]
282    #[buildable(root = crate)]
283    #[builder(root = crate)]
284    pub struct MseLoss {
285        /// Reduction type. Default to [`MseLoss::DEFAULT_REDUCTION`]
286        #[builder(optional, default = MseLoss::DEFAULT_REDUCTION)]
287        pub reduction: LossReduction,
288    }
289}
290
291impl MseLoss {
292    /// Default value for the reduction parameter.
293    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::Mean;
294
295    /// Compute the mean squared error loss.
296    ///
297    /// # Params
298    ///
299    /// - `predictions`: predicted values
300    /// - `targets`: target values
301    pub fn apply(
302        &self,
303        predictions: impl AsRef<Array>,
304        targets: impl AsRef<Array>,
305    ) -> Result<Array, Exception> {
306        let predictions = predictions.as_ref();
307        let targets = targets.as_ref();
308        let reduction = self.reduction;
309
310        check_shape(predictions, targets, "predictions", "targets")?;
311        let loss = predictions.subtract(targets)?.square()?;
312        reduction.reduce(loss)
313    }
314}
315
316generate_builder! {
317    /// Computes the negative log likelihood loss.
318    #[derive(Debug, Clone, Buildable)]
319    #[buildable(root = crate)]
320    #[builder(root = crate)]
321    pub struct NllLoss {
322        /// distribution axis. Default to [`NllLoss::DEFAULT_AXIS`]
323        #[builder(optional, default = NllLoss::DEFAULT_AXIS)]
324        pub axis: i32,
325
326        /// Reduction type. Default to [`NllLoss::DEFAULT_REDUCTION`]
327        #[builder(optional, default = NllLoss::DEFAULT_REDUCTION)]
328        pub reduction: LossReduction,
329    }
330}
331
332impl NllLoss {
333    /// Default value for the `axis` parameter.
334    pub const DEFAULT_AXIS: i32 = -1;
335
336    /// Default value for the `reduction` parameter.
337    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
338
339    /// Compute the negative log likelihood loss.
340    ///
341    /// # Params
342    ///
343    /// - `inputs`: predicted distribution in log space
344    /// - `targets`: target values
345    pub fn apply(
346        &self,
347        inputs: impl AsRef<Array>,
348        targets: impl AsRef<Array>,
349    ) -> Result<Array, Exception> {
350        let inputs = inputs.as_ref();
351        let targets = targets.as_ref();
352        let axis = self.axis;
353        let reduction = self.reduction;
354
355        let loss = -take_along_axis(inputs, &targets.expand_dims_axes(&[-1])?, axis)?
356            .squeeze_axes(&[-1])?;
357        reduction.reduce(loss)
358    }
359}
360
361generate_builder! {
362    /// Compute the negative log likelihood loss for a Gaussian distribution.
363    #[derive(Debug, Clone, Buildable)]
364    #[buildable(root = crate)]
365    #[builder(root = crate)]
366    pub struct GaussianNllLoss {
367        /// Whether to include the constant term in the loss calculation. Default to
368        /// [`GaussianNllLoss::DEFAULT_FULL`]
369        #[builder(optional, default = GaussianNllLoss::DEFAULT_FULL)]
370        pub full: bool,
371
372        /// Small positive constant for numerical stability. Default to
373        /// [`GaussianNllLoss::DEFAULT_EPS`]
374        #[builder(optional, default = GaussianNllLoss::DEFAULT_EPS)]
375        pub eps: f32,
376
377        /// Reduction type. Default to [`GaussianNllLoss::DEFAULT_REDUCTION`]
378        #[builder(optional, default = GaussianNllLoss::DEFAULT_REDUCTION)]
379        pub reduction: LossReduction,
380    }
381}
382
383impl GaussianNllLoss {
384    /// Default value for the `full` parameter.
385    pub const DEFAULT_FULL: bool = false;
386
387    /// Default value for the `eps` parameter.
388    pub const DEFAULT_EPS: f32 = 1e-6;
389
390    /// Default value for the `reduction` parameter.
391    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
392
393    /// Compute the negative log likelihood loss for a Gaussian distribution.
394    ///
395    /// # Params
396    ///
397    /// - `inputs`: The predicted expectation of the Gaussian distribution.
398    /// - `targets`: The target values (samples from the Gaussian distribution).
399    /// - `vars`: The predicted variance of the Gaussian distribution.
400    pub fn apply(
401        &self,
402        inputs: impl AsRef<Array>,
403        targets: impl AsRef<Array>,
404        vars: impl AsRef<Array>,
405    ) -> Result<Array, Exception> {
406        let inputs = inputs.as_ref();
407        let targets = targets.as_ref();
408        let vars = vars.as_ref();
409        let full = self.full;
410        let eps = self.eps;
411        let reduction = self.reduction;
412
413        check_shape(inputs, targets, "inputs", "targets")?;
414        check_shape(inputs, vars, "inputs", "vars")?;
415
416        let vars = maximum(vars, array!(eps))?;
417        let mut loss =
418            array!(0.5) * (log(&vars)?.add(square(&targets.subtract(inputs)?)?.divide(&vars)?)?);
419
420        if full {
421            let pi = array!(std::f32::consts::PI);
422            loss = loss.add(array!(0.5).multiply(log(&array!(2.0).multiply(pi)?)?)?)?;
423        }
424
425        reduction.reduce(loss)
426    }
427}
428
429generate_builder! {
430    /// Compute the Kullback-Leibler divergence loss.
431    ///
432    /// Computes the following when the `reduction` is `LossReduction::None`:
433    ///
434    /// ```rust, ignore
435    /// sum(exp(targets) * (targets - inputs), axis, None)
436    /// ```
437    #[derive(Debug, Clone, Buildable)]
438    #[buildable(root = crate)]
439    #[builder(root = crate)]
440    pub struct KlDivLoss {
441        /// The distribution axis. Default to [`KlDivLoss::DEFAULT_AXIS`]
442        #[builder(optional, default = KlDivLoss::DEFAULT_AXIS)]
443        pub axis: i32,
444
445        /// Reduction type. Default to [`KlDivLoss::DEFAULT_REDUCTION`]
446        #[builder(optional, default = KlDivLoss::DEFAULT_REDUCTION)]
447        pub reduction: LossReduction,
448    }
449}
450
451impl KlDivLoss {
452    /// Default value for the `axis` parameter.
453    pub const DEFAULT_AXIS: i32 = -1;
454
455    /// Default value for the `reduction` parameter.
456    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
457
458    /// Compute the Kullback-Leibler divergence loss.
459    ///
460    /// # Params
461    ///
462    /// - `inputs`: Log probabilities for the predicted distribution.
463    /// - `targets`: Log probabilities for the target distribution.
464    pub fn apply(
465        &self,
466        inputs: impl AsRef<Array>,
467        targets: impl AsRef<Array>,
468    ) -> Result<Array, Exception> {
469        let inputs = inputs.as_ref();
470        let targets = targets.as_ref();
471        let axis = self.axis;
472        let reduction = self.reduction;
473
474        let loss = sum_axis(
475            &exp(targets)?.multiply(targets.subtract(inputs)?)?,
476            axis,
477            None,
478        )?;
479        reduction.reduce(loss)
480    }
481}
482
483generate_builder! {
484    /// Computes the smooth L1 loss.
485    ///
486    /// The smooth L1 loss is a variant of the L1 loss which replaces the absolute
487    /// difference with a squared difference when the absolute difference is less
488    /// than `beta`.
489    #[derive(Debug, Clone, Buildable)]
490    #[buildable(root = crate)]
491    #[builder(root = crate)]
492    pub struct SmoothL1Loss {
493        /// The threshold after which the loss changes from the squared to the absolute difference.
494        /// Default to [`SmoothL1Loss::DEFAULT_BETA`]
495        #[builder(optional, default = SmoothL1Loss::DEFAULT_BETA)]
496        pub beta: f32,
497
498        /// Reduction type. Default to [`SmoothL1Loss::DEFAULT_REDUCTION`]
499        #[builder(optional, default = SmoothL1Loss::DEFAULT_REDUCTION)]
500        pub reduction: LossReduction,
501    }
502}
503
504impl SmoothL1Loss {
505    /// Default value for the `beta` parameter.
506    pub const DEFAULT_BETA: f32 = 1.0;
507
508    /// Default value for the `reduction` parameter.
509    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::Mean;
510
511    /// Compute the smooth L1 loss.
512    ///
513    /// # Params
514    ///
515    /// - `predictions`: predicted values
516    /// - `targets`: target values
517    pub fn apply(
518        &self,
519        predictions: impl AsRef<Array>,
520        targets: impl AsRef<Array>,
521    ) -> Result<Array, Exception> {
522        let predictions = predictions.as_ref();
523        let targets = targets.as_ref();
524        let beta = self.beta;
525        let reduction = self.reduction;
526
527        check_shape(predictions, targets, "predictions", "targets")?;
528        let diff = predictions.subtract(targets)?.abs()?;
529        let beta = array!(beta);
530        let loss = select(
531            &diff.lt(&beta)?,
532            array!(0.5).multiply(square(&diff)?)?.divide(&beta)?,
533            diff.subtract(array!(0.5).multiply(beta)?)?,
534        )?;
535        reduction.reduce(loss)
536    }
537}
538
539generate_builder! {
540    /// Computes the triplet loss for a set of anchor, positive, and negative samples. Margin is
541    /// represented with alpha in the math section.
542    #[derive(Debug, Clone, Buildable)]
543    #[buildable(root = crate)]
544    #[builder(root = crate)]
545    pub struct TripletLoss {
546        /// Distribution axis. Default to [`TripletLoss::DEFAULT_AXIS`]
547        #[builder(optional, default = TripletLoss::DEFAULT_AXIS)]
548        pub axis: i32,
549
550        /// The norm degree for pairwise distance. Default to [`TripletLoss::DEFAULT_P`]
551        #[builder(optional, default = TripletLoss::DEFAULT_P)]
552        pub p: f32,
553
554        /// Margin for the triplet loss. Default to [`TripletLoss::DEFAULT_MARGIN`]
555        #[builder(optional, default = TripletLoss::DEFAULT_MARGIN)]
556        pub margin: f32,
557
558        /// Small positive constant for numerical stability. Default to [`TripletLoss::DEFAULT_EPS`]
559        #[builder(optional, default = TripletLoss::DEFAULT_EPS)]
560        pub eps: f32,
561
562        /// Reduction type. Default to [`TripletLoss::DEFAULT_REDUCTION`]
563        #[builder(optional, default = TripletLoss::DEFAULT_REDUCTION)]
564        pub reduction: LossReduction,
565    }
566}
567
568impl TripletLoss {
569    /// Default value for the `axis` parameter.
570    pub const DEFAULT_AXIS: i32 = -1;
571
572    /// Default value for the `p` parameter.
573    pub const DEFAULT_P: f32 = 2.0;
574
575    /// Default value for the `margin` parameter.
576    pub const DEFAULT_MARGIN: f32 = 1.0;
577
578    /// Default value for the `eps` parameter.
579    pub const DEFAULT_EPS: f32 = 1e-6;
580
581    /// Default value for the `reduction` parameter.
582    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
583
584    /// Computes the triplet loss for a set of anchor, positive, and negative samples. Margin is
585    /// represented with alpha in the math section.
586    ///
587    /// # Params
588    ///
589    /// - `anchors`: The anchor samples
590    /// - `positives`: The positive samples
591    /// - `neonatives`: The negative samples
592    pub fn apply(
593        &self,
594        anchors: impl AsRef<Array>,
595        positives: impl AsRef<Array>,
596        negatives: impl AsRef<Array>,
597    ) -> Result<Array, Exception> {
598        let anchors = anchors.as_ref();
599        let positives = positives.as_ref();
600        let negatives = negatives.as_ref();
601        let axis = self.axis;
602        let p = self.p;
603        let margin = self.margin;
604        let eps = self.eps;
605        let reduction = self.reduction;
606
607        let eps = array!(eps);
608        let p = array!(p);
609        let margin = array!(margin);
610
611        let pos = sqrt(
612            &power(&anchors.subtract(positives)?, &p)?
613                .sum_axis(axis, None)?
614                .add(&eps)?,
615        )?;
616        let neg = sqrt(
617            &power(&anchors.subtract(negatives)?, &p)?
618                .sum_axis(axis, None)?
619                .add(&eps)?,
620        )?;
621        let loss = maximum(pos.subtract(neg)?.add(margin)?, array!(0.0))?;
622        reduction.reduce(loss)
623    }
624}
625
626generate_builder! {
627    /// Compute the hinge loss.
628    #[derive(Debug, Clone, Buildable)]
629    #[buildable(root = crate)]
630    #[builder(root = crate)]
631    pub struct HingeLoss {
632        /// Reduction type. Default to [`HingeLoss::DEFAULT_REDUCTION`]
633        #[builder(optional, default = HingeLoss::DEFAULT_REDUCTION)]
634        pub reduction: LossReduction,
635    }
636}
637
638impl HingeLoss {
639    /// Default value for the `reduction` parameter.
640    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
641
642    /// Compute the hinge loss.
643    ///
644    /// # Params
645    ///
646    /// - `inputs`: predicted values
647    /// - `targets`: target values, -1 or 1
648    pub fn apply(
649        &self,
650        inputs: impl AsRef<Array>,
651        targets: impl AsRef<Array>,
652    ) -> Result<Array, Exception> {
653        let inputs = inputs.as_ref();
654        let targets = targets.as_ref();
655        let reduction = self.reduction;
656
657        let a = array!(1.0).subtract(inputs.multiply(targets)?)?;
658        let b = array!(0.0);
659        let loss = maximum(a, b)?;
660        reduction.reduce(loss)
661    }
662}
663
664generate_builder! {
665    /// Compute the Huber loss.
666    #[derive(Debug, Clone, Buildable)]
667    #[buildable(root = crate)]
668    #[builder(root = crate)]
669    pub struct HuberLoss {
670        /// The threshold at which to change between L1 and L2 loss. Default to
671        /// [`HuberLoss::DEFAULT_DELTA`]
672        #[builder(optional, default = HuberLoss::DEFAULT_DELTA)]
673        pub delta: f32,
674
675        /// Reduction type. Default to [`HuberLoss::DEFAULT_REDUCTION`]
676        #[builder(optional, default = HuberLoss::DEFAULT_REDUCTION)]
677        pub reduction: LossReduction,
678    }
679}
680
681impl HuberLoss {
682    /// Default value for the `delta` parameter.
683    pub const DEFAULT_DELTA: f32 = 1.0;
684
685    /// Default value for the `reduction` parameter.
686    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
687
688    /// Compute the Huber loss.
689    ///
690    /// # Params
691    ///
692    /// - `inputs`: predicted values
693    /// - `targets`: target values
694    pub fn apply(
695        &self,
696        inputs: impl AsRef<Array>,
697        targets: impl AsRef<Array>,
698    ) -> Result<Array, Exception> {
699        let inputs = inputs.as_ref();
700        let targets = targets.as_ref();
701        let delta = self.delta;
702        let reduction = self.reduction;
703
704        let errors = inputs.subtract(targets)?;
705        let abs_errors = errors.abs()?;
706        let quadratic = minimum(&abs_errors, array!(delta))?;
707        let linear = abs_errors.subtract(&quadratic)?;
708        let loss = array!(0.5)
709            .multiply(square(&quadratic)?)?
710            .add(array!(delta).multiply(linear)?)?;
711        reduction.reduce(loss)
712    }
713}
714
715generate_builder! {
716    /// Computes the log cosh loss between inputs and targets.
717    ///
718    /// Logcosh acts like L2 loss for small errors, ensuring stable gradients,
719    /// and like the L1 loss for large errors, reducing sensitivity to outliers. This
720    /// dual behavior offers a balanced, robust approach for regression tasks.
721    #[derive(Debug, Clone, Buildable)]
722    #[buildable(root = crate)]
723    #[builder(root = crate)]
724    pub struct LogCoshLoss {
725        /// Reduction type. Default to [`LogCoshLoss::DEFAULT_REDUCTION`]
726        #[builder(optional, default = LogCoshLoss::DEFAULT_REDUCTION)]
727        pub reduction: LossReduction,
728    }
729}
730
731impl LogCoshLoss {
732    /// Default value for the `reduction` parameter.
733    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
734
735    /// Computes the log cosh loss between inputs and targets.
736    ///
737    /// # Params
738    ///
739    /// - `inputs`: predicted values
740    /// - `targets`: target values
741    pub fn apply(
742        &self,
743        inputs: impl AsRef<Array>,
744        targets: impl AsRef<Array>,
745    ) -> Result<Array, Exception> {
746        let inputs = inputs.as_ref();
747        let targets = targets.as_ref();
748        let reduction = self.reduction;
749
750        let errors = inputs.subtract(targets)?;
751        let neg_errors = errors.negative()?;
752        let loss = logaddexp(errors, neg_errors)?.subtract(log(&array!(2.0))?)?;
753        reduction.reduce(loss)
754    }
755}
756
757generate_builder! {
758    /// Computes the cosine similarity loss.
759    #[derive(Debug, Clone, Buildable)]
760    #[buildable(root = crate)]
761    #[builder(root = crate)]
762    pub struct CosineSimilarityLoss {
763        /// Embedding axis. Default to [`CosineSimilarityLoss::DEFAULT_AXIS`]
764        #[builder(optional, default = CosineSimilarityLoss::DEFAULT_AXIS)]
765        pub axis: i32,
766
767        /// minimum value of the denominator used for numerical stability. Default to
768        /// [`CosineSimilarityLoss::DEFAULT_EPS`]
769        #[builder(optional, default = CosineSimilarityLoss::DEFAULT_EPS)]
770        pub eps: f32,
771
772        /// Reduction type. Default to [`CosineSimilarityLoss::DEFAULT_REDUCTION`]
773        #[builder(optional, default = CosineSimilarityLoss::DEFAULT_REDUCTION)]
774        pub reduction: LossReduction,
775    }
776}
777
778impl CosineSimilarityLoss {
779    /// Default value for the `axis` parameter.
780    pub const DEFAULT_AXIS: i32 = -1;
781
782    /// Default value for the `eps` parameter.
783    pub const DEFAULT_EPS: f32 = 1e-8;
784
785    /// Default value for the `reduction` parameter.
786    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
787
788    /// Computes the cosine similarity loss.
789    ///
790    /// # Params
791    ///
792    /// - `x1`: first array
793    /// - `x2`: second array
794    pub fn apply(&self, x1: impl AsRef<Array>, x2: impl AsRef<Array>) -> Result<Array, Exception> {
795        let x1 = x1.as_ref();
796        let x2 = x2.as_ref();
797        let axis = self.axis;
798        let eps = self.eps;
799        let reduction = self.reduction;
800
801        fn l2_loss(a: &Array, axis: i32) -> Result<Array, Exception> {
802            if a.dtype().is_complex() {
803                Ok(sqrt(&sum_axis(&abs(a)?.square()?, axis, None)?)?)
804            } else {
805                Ok(sqrt(&sum_axis(&a.square()?, axis, None)?)?)
806            }
807        }
808
809        let x1_norm = l2_loss(x1, axis)?;
810        let x2_norm = l2_loss(x2, axis)?;
811
812        let num = sum_axis(&x1.multiply(x2)?, axis, None)?;
813        let den = maximum(x1_norm.multiply(x2_norm)?, array!(eps))?;
814        let loss = num.divide(&den)?;
815
816        reduction.reduce(loss)
817    }
818}
819
820generate_builder! {
821    /// Computes the margin ranking loss.
822    #[derive(Debug, Clone, Buildable)]
823    #[buildable(root = crate)]
824    #[builder(root = crate)]
825    pub struct MarginRankingLoss {
826        /// The margin by which the scores should be separated. Default to
827        /// [`MarginRankingLoss::DEFAULT_MARGIN`]
828        #[builder(optional, default = MarginRankingLoss::DEFAULT_MARGIN)]
829        pub margin: f32,
830
831        /// Reduction type. Default to [`MarginRankingLoss::DEFAULT_REDUCTION`]
832        #[builder(optional, default = MarginRankingLoss::DEFAULT_REDUCTION)]
833        pub reduction: LossReduction,
834    }
835}
836
837impl MarginRankingLoss {
838    /// Default value for the `margin` parameter.
839    pub const DEFAULT_MARGIN: f32 = 0.0;
840
841    /// Default value for the `reduction` parameter.
842    pub const DEFAULT_REDUCTION: LossReduction = LossReduction::None;
843
844    /// Computes the margin ranking loss.
845    ///
846    /// # Params
847    ///
848    /// - `inputs1`: Scores for the first input.
849    /// - `inputs2`: Scores for the second input.
850    /// - `targets`: Labels indicating whether samples in `inputs1` should be ranked higher than samples
851    ///   in `inputs2`. Values should be 1 or -1.
852    pub fn apply(
853        &self,
854        inputs1: impl AsRef<Array>,
855        inputs2: impl AsRef<Array>,
856        targets: impl AsRef<Array>,
857    ) -> Result<Array, Exception> {
858        let inputs1 = inputs1.as_ref();
859        let inputs2 = inputs2.as_ref();
860        let targets = targets.as_ref();
861        let margin = self.margin;
862        let reduction = self.reduction;
863
864        check_shape(inputs1, inputs2, "inputs1", "inputs2")?;
865        check_shape(inputs1, targets, "inputs1", "targets")?;
866
867        let margin = array!(margin);
868        let diff = inputs1.subtract(inputs2)?;
869        let loss = maximum(
870            array!(0.0),
871            targets.multiply(diff)?.negative()?.add(margin)?,
872        )?;
873        reduction.reduce(loss)
874    }
875}
876
877#[cfg(test)]
878#[allow(clippy::approx_constant)]
879mod tests {
880    use crate::{
881        array,
882        builder::Builder,
883        ops::is_nan,
884        test_utils::{assert_array_eq, tolerances},
885    };
886    use float_eq::assert_float_eq;
887
888    use super::*;
889
890    // The following unit tests are adapted from the python API at: mlx/python/tests/test_losses.py
891
892    #[test]
893    fn test_cross_entropy() {
894        // No weights, no label smoothing
895        let logits = array!([[0.0, f32::NEG_INFINITY], [f32::NEG_INFINITY, 0.0]]);
896        let indices = array!([0, 1]);
897        let expected = array!([0.0, 0.0]);
898        let loss = CrossEntropy::new()
899            .unwrap()
900            .apply(&logits, indices)
901            .unwrap();
902        assert_array_eq(
903            loss,
904            &expected,
905            tolerances::MLX_DEFAULT.rtol,
906            tolerances::MLX_DEFAULT.atol,
907        );
908
909        let probs = array!([[1.0, 0.0], [0.0, 1.0]]);
910        let cross_entropy = CrossEntropyBuilder::new()
911            .reduction(LossReduction::None)
912            .build()
913            .unwrap();
914        let loss = cross_entropy.apply(logits, probs).unwrap();
915        assert!(is_nan(&loss)
916            .unwrap()
917            .all(None)
918            .unwrap()
919            .item_exact::<bool>());
920
921        // With weights, no label smoothing
922        let logits = array!([[2.0, -1.0], [-1.0, 2.0]]);
923        let indices = array!([0, 1]);
924        let weights = array!([1.0, 2.0]);
925        let expected = array!([0.04858735, 0.0971747]);
926        let cross_entropy = CrossEntropyBuilder::new()
927            .weights(&weights)
928            .reduction(LossReduction::None)
929            .build()
930            .unwrap();
931        let loss = cross_entropy.apply(&logits, indices).unwrap();
932        assert_array_eq(
933            loss,
934            &expected,
935            tolerances::MLX_DEFAULT.rtol,
936            tolerances::MLX_DEFAULT.atol,
937        );
938
939        let probs = array!([[1.0, 0.0], [0.0, 1.0]]);
940        let cross_entropy = CrossEntropyBuilder::new()
941            .weights(&weights)
942            .reduction(LossReduction::None)
943            .build()
944            .unwrap();
945        let loss = cross_entropy.apply(logits, probs).unwrap();
946        assert_array_eq(
947            loss,
948            &expected,
949            tolerances::MLX_DEFAULT.rtol,
950            tolerances::MLX_DEFAULT.atol,
951        );
952
953        // No weights, with label smoothing
954        let logits = array!([[2.0, -1.0], [-1.0, 2.0]]);
955        let indices = array!([0, 1]);
956        let expected = array!([0.498587, 0.498587]);
957        let cross_entropy = CrossEntropyBuilder::new()
958            .label_smoothing(0.3)
959            .reduction(LossReduction::None)
960            .build()
961            .unwrap();
962        let loss = cross_entropy.apply(&logits, indices).unwrap();
963        assert_array_eq(
964            loss,
965            &expected,
966            tolerances::MLX_DEFAULT.rtol,
967            tolerances::MLX_DEFAULT.atol,
968        );
969
970        let probs = array!([[1.0, 0.0], [0.0, 1.0]]);
971        let cross_entropy = CrossEntropyBuilder::new()
972            .label_smoothing(0.3)
973            .reduction(LossReduction::None)
974            .build()
975            .unwrap();
976        let loss = cross_entropy.apply(logits, probs).unwrap();
977        assert_array_eq(
978            loss,
979            expected,
980            tolerances::MLX_DEFAULT.rtol,
981            tolerances::MLX_DEFAULT.atol,
982        );
983
984        // With weights and label smoothing
985        let logits = array!([[2.0, -1.0], [-1.0, 2.0]]);
986        let indices = array!([0, 1]);
987        let weights = array!([1.0, 2.0]);
988        let expected = array!([0.49858734, 0.9971747]);
989        let cross_entropy = CrossEntropyBuilder::new()
990            .weights(&weights)
991            .label_smoothing(0.3)
992            .reduction(LossReduction::None)
993            .build()
994            .unwrap();
995        let loss = cross_entropy.apply(&logits, indices).unwrap();
996        assert_array_eq(
997            loss,
998            &expected,
999            tolerances::MLX_DEFAULT.rtol,
1000            tolerances::MLX_DEFAULT.atol,
1001        );
1002
1003        let probs = array!([[1.0, 0.0], [0.0, 1.0]]);
1004        let cross_entropy = CrossEntropyBuilder::new()
1005            .weights(&weights)
1006            .label_smoothing(0.3)
1007            .reduction(LossReduction::None)
1008            .build()
1009            .unwrap();
1010        let loss = cross_entropy.apply(logits, probs).unwrap();
1011        assert_array_eq(
1012            loss,
1013            expected,
1014            tolerances::MLX_DEFAULT.rtol,
1015            tolerances::MLX_DEFAULT.atol,
1016        );
1017    }
1018
1019    #[test]
1020    fn test_binary_cross_entropy_with_logits_as_inputs() {
1021        let logits = array!([0.105361, 0.223144, 1.20397, 0.916291]);
1022        let targets = array!([0.0, 0.0, 1.0, 1.0]);
1023
1024        // Test with reduction 'none'
1025        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1026            .reduction(LossReduction::None)
1027            .build()
1028            .unwrap();
1029        let loss_none = binary_cross_entropy.apply(&logits, &targets).unwrap();
1030        let expected_none = array!([0.747215, 0.810930, 0.262365, 0.336472]);
1031        assert_array_eq(
1032            loss_none,
1033            &expected_none,
1034            tolerances::MLX_DEFAULT.rtol,
1035            tolerances::MLX_DEFAULT.atol,
1036        );
1037
1038        // Test with reduction 'mean'
1039        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1040            .reduction(LossReduction::Mean)
1041            .build()
1042            .unwrap();
1043        let loss_mean = binary_cross_entropy.apply(&logits, &targets).unwrap();
1044        let expected_mean = expected_none.mean(None).unwrap();
1045        assert_array_eq(
1046            loss_mean,
1047            expected_mean,
1048            tolerances::MLX_DEFAULT.rtol,
1049            tolerances::MLX_DEFAULT.atol,
1050        );
1051
1052        // Test with reduction 'sum'
1053        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1054            .reduction(LossReduction::Sum)
1055            .build()
1056            .unwrap();
1057        let loss = binary_cross_entropy.apply(&logits, &targets).unwrap();
1058        let expected = expected_none.sum(None).unwrap();
1059        assert_array_eq(
1060            loss,
1061            expected,
1062            tolerances::MLX_DEFAULT.rtol,
1063            tolerances::MLX_DEFAULT.atol,
1064        );
1065
1066        // With weights, no label smoothing
1067        let weights = array!([1.0, 2.0, 1.0, 2.0]);
1068        let expected = array!([0.747215, 1.62186, 0.262365, 0.672944]);
1069        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1070            .weights(&weights)
1071            .reduction(LossReduction::None)
1072            .build()
1073            .unwrap();
1074        let loss = binary_cross_entropy.apply(&logits, &targets).unwrap();
1075        assert_array_eq(
1076            loss,
1077            expected,
1078            tolerances::MLX_DEFAULT.rtol,
1079            tolerances::MLX_DEFAULT.atol,
1080        );
1081    }
1082
1083    #[test]
1084    fn test_binary_cross_entropy_with_probs_as_inputs() {
1085        let probs = array!([0.5, 0.6, 0.7, 0.8]);
1086        let targets = array!([0.0, 0.0, 1.0, 1.0]);
1087
1088        // Test with reduction 'none'
1089        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1090            .inputs_are_logits(false)
1091            .reduction(LossReduction::None)
1092            .build()
1093            .unwrap();
1094        let loss_none = binary_cross_entropy.apply(&probs, &targets).unwrap();
1095        let expected_none = array!([0.693147, 0.916291, 0.356675, 0.223144]);
1096        assert_array_eq(
1097            loss_none,
1098            &expected_none,
1099            tolerances::MLX_DEFAULT.rtol,
1100            tolerances::MLX_DEFAULT.atol,
1101        );
1102
1103        // Test with reduction 'mean'
1104        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1105            .inputs_are_logits(false)
1106            .reduction(LossReduction::Mean)
1107            .build()
1108            .unwrap();
1109        let loss_mean = binary_cross_entropy.apply(&probs, &targets).unwrap();
1110        let expected_mean = expected_none.mean(None).unwrap();
1111        assert_array_eq(
1112            loss_mean,
1113            expected_mean,
1114            tolerances::MLX_DEFAULT.rtol,
1115            tolerances::MLX_DEFAULT.atol,
1116        );
1117
1118        // Test with reduction 'sum'
1119        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1120            .inputs_are_logits(false)
1121            .reduction(LossReduction::Sum)
1122            .build()
1123            .unwrap();
1124        let loss = binary_cross_entropy.apply(&probs, &targets).unwrap();
1125        let expected = expected_none.sum(None).unwrap();
1126        assert_array_eq(
1127            loss,
1128            expected,
1129            tolerances::MLX_DEFAULT.rtol,
1130            tolerances::MLX_DEFAULT.atol,
1131        );
1132    }
1133
1134    #[test]
1135    fn test_binary_cross_entropy_with_tiny_probs_as_inputs() {
1136        let tiny_prob = 1e-59;
1137        let probs = array!([0.0, tiny_prob, 1.0 - tiny_prob, 1.0]);
1138        let targets = array!([0.0, 0.0, 1.0, 1.0]);
1139
1140        // Test with reduction 'none'
1141        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1142            .inputs_are_logits(false)
1143            .reduction(LossReduction::None)
1144            .build()
1145            .unwrap();
1146        let loss_none = binary_cross_entropy.apply(&probs, &targets).unwrap();
1147        let expected_none = array!([0.0, tiny_prob, tiny_prob, 0.0]);
1148        assert_array_eq(
1149            loss_none,
1150            &expected_none,
1151            tolerances::MLX_DEFAULT.rtol,
1152            tolerances::MLX_DEFAULT.atol,
1153        );
1154
1155        // Test with reduction 'mean'
1156        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1157            .inputs_are_logits(false)
1158            .reduction(LossReduction::Mean)
1159            .build()
1160            .unwrap();
1161        let loss_mean = binary_cross_entropy.apply(&probs, &targets).unwrap();
1162        let expected_mean = expected_none.mean(None).unwrap();
1163        assert_array_eq(
1164            loss_mean,
1165            expected_mean,
1166            tolerances::MLX_DEFAULT.rtol,
1167            tolerances::MLX_DEFAULT.atol,
1168        );
1169
1170        // Test with reduction 'sum'
1171        let binary_cross_entropy = BinaryCrossEntropyBuilder::new()
1172            .inputs_are_logits(false)
1173            .reduction(LossReduction::Sum)
1174            .build()
1175            .unwrap();
1176        let loss = binary_cross_entropy.apply(&probs, &targets).unwrap();
1177        let expected = expected_none.sum(None).unwrap();
1178        assert_array_eq(
1179            loss,
1180            expected,
1181            tolerances::MLX_DEFAULT.rtol,
1182            tolerances::MLX_DEFAULT.atol,
1183        );
1184    }
1185
1186    #[test]
1187    fn test_l1_loss() {
1188        let predictions = array!([0.5, 0.2, 0.9, 0.0]);
1189        let targets = array!([0.5, 0.2, 0.9, 0.0]);
1190
1191        let expected_none = array!([0.0, 0.0, 0.0, 0.0]);
1192        let expected_sum = expected_none.sum(None).unwrap();
1193        let expected_mean = expected_none.mean(None).unwrap();
1194
1195        let l1_loss = L1LossBuilder::new()
1196            .reduction(LossReduction::None)
1197            .build()
1198            .unwrap();
1199        let loss_none = l1_loss.apply(&predictions, &targets).unwrap();
1200        assert_array_eq(
1201            loss_none,
1202            &expected_none,
1203            tolerances::MLX_DEFAULT.rtol,
1204            tolerances::MLX_DEFAULT.atol,
1205        );
1206
1207        let l1_loss = L1LossBuilder::new()
1208            .reduction(LossReduction::Sum)
1209            .build()
1210            .unwrap();
1211        let loss_sum = l1_loss.apply(&predictions, &targets).unwrap();
1212        assert_array_eq(
1213            loss_sum,
1214            expected_sum,
1215            tolerances::MLX_DEFAULT.rtol,
1216            tolerances::MLX_DEFAULT.atol,
1217        );
1218
1219        let l1_loss = L1LossBuilder::new()
1220            .reduction(LossReduction::Mean)
1221            .build()
1222            .unwrap();
1223        let loss_mean = l1_loss.apply(&predictions, &targets).unwrap();
1224        assert_array_eq(
1225            loss_mean,
1226            expected_mean,
1227            tolerances::MLX_DEFAULT.rtol,
1228            tolerances::MLX_DEFAULT.atol,
1229        );
1230    }
1231
1232    #[test]
1233    fn test_mse_loss() {
1234        let predictions = array!([0.5, 0.2, 0.9, 0.0]);
1235        let targets = array!([0.7, 0.1, 0.8, 0.2]);
1236
1237        let expected_none = array!([0.04, 0.01, 0.01, 0.04]);
1238        let expected_mean = expected_none.mean(None).unwrap();
1239        let expected_sum = expected_none.sum(None).unwrap();
1240
1241        let mse_loss = MseLossBuilder::new()
1242            .reduction(LossReduction::None)
1243            .build()
1244            .unwrap();
1245        let loss_none = mse_loss.apply(&predictions, &targets).unwrap();
1246        assert_array_eq(
1247            loss_none,
1248            expected_none,
1249            tolerances::MLX_DEFAULT.rtol,
1250            tolerances::MLX_DEFAULT.atol,
1251        );
1252
1253        let mse_loss = MseLossBuilder::new()
1254            .reduction(LossReduction::Mean)
1255            .build()
1256            .unwrap();
1257        let loss_mean = mse_loss.apply(&predictions, &targets).unwrap();
1258        assert_array_eq(
1259            loss_mean,
1260            expected_mean,
1261            tolerances::MLX_DEFAULT.rtol,
1262            tolerances::MLX_DEFAULT.atol,
1263        );
1264
1265        let mse_loss = MseLossBuilder::new()
1266            .reduction(LossReduction::Sum)
1267            .build()
1268            .unwrap();
1269        let loss_sum = mse_loss.apply(&predictions, &targets).unwrap();
1270        assert_array_eq(
1271            loss_sum,
1272            expected_sum,
1273            tolerances::MLX_DEFAULT.rtol,
1274            tolerances::MLX_DEFAULT.atol,
1275        );
1276    }
1277
1278    #[test]
1279    fn test_smooth_l1_loss() {
1280        let predictions = array!([1.5, 2.5, 0.5, 3.5]);
1281        let targets = array!([1.0, 2.0, 0.5, 2.5]);
1282        let beta = 1.0;
1283
1284        let expected_none = array!([0.125, 0.125, 0.0, 0.5]);
1285        let expected_sum = expected_none.sum(None).unwrap();
1286        let expected_mean = expected_none.mean(None).unwrap();
1287
1288        let smooth_l1_loss = SmoothL1LossBuilder::new()
1289            .beta(beta)
1290            .reduction(LossReduction::None)
1291            .build()
1292            .unwrap();
1293        let loss_none = smooth_l1_loss.apply(&predictions, &targets).unwrap();
1294        assert_array_eq(
1295            loss_none,
1296            expected_none,
1297            tolerances::MLX_DEFAULT.rtol,
1298            tolerances::MLX_DEFAULT.atol,
1299        );
1300
1301        let smooth_l1_loss = SmoothL1LossBuilder::new()
1302            .beta(beta)
1303            .reduction(LossReduction::Sum)
1304            .build()
1305            .unwrap();
1306        let loss_sum = smooth_l1_loss.apply(&predictions, &targets).unwrap();
1307        assert_array_eq(
1308            loss_sum,
1309            expected_sum,
1310            tolerances::MLX_DEFAULT.rtol,
1311            tolerances::MLX_DEFAULT.atol,
1312        );
1313
1314        let smooth_l1_loss = SmoothL1LossBuilder::new()
1315            .beta(beta)
1316            .reduction(LossReduction::Mean)
1317            .build()
1318            .unwrap();
1319        let loss_mean = smooth_l1_loss.apply(&predictions, &targets).unwrap();
1320        assert_array_eq(
1321            loss_mean,
1322            expected_mean,
1323            tolerances::MLX_DEFAULT.rtol,
1324            tolerances::MLX_DEFAULT.atol,
1325        );
1326    }
1327
1328    #[test]
1329    fn test_smooth_l1_loss_negative_diff() {
1330        let a = array!([1.5, 6.0, 0.5, 2.5]);
1331        let b = array!([1.0, 2.0, 0.5, 3.5]);
1332
1333        let loss = SmoothL1Loss::new();
1334
1335        let ab = loss.apply(&a, &b).unwrap();
1336        let ba = loss.apply(&b, &a).unwrap();
1337        assert_array_eq(
1338            ab,
1339            ba,
1340            tolerances::MLX_DEFAULT.rtol,
1341            tolerances::MLX_DEFAULT.atol,
1342        );
1343    }
1344
1345    #[test]
1346    fn test_nll_loss() {
1347        let logits = array!([[0.0, f32::NEG_INFINITY], [f32::NEG_INFINITY, 0.0]]);
1348        let targets = array!([0, 1]);
1349
1350        let expected_none = array!([0.0, 0.0]);
1351        let expected_sum = expected_none.sum(None).unwrap();
1352        let expected_mean = expected_none.mean(None).unwrap();
1353
1354        let nll_loss = NllLossBuilder::new()
1355            .reduction(LossReduction::None)
1356            .build()
1357            .unwrap();
1358        let loss_none = nll_loss.apply(&logits, &targets).unwrap();
1359        assert_array_eq(
1360            loss_none,
1361            expected_none,
1362            tolerances::MLX_DEFAULT.rtol,
1363            tolerances::MLX_DEFAULT.atol,
1364        );
1365
1366        let nll_loss = NllLossBuilder::new()
1367            .reduction(LossReduction::Mean)
1368            .build()
1369            .unwrap();
1370        let loss_mean = nll_loss.apply(&logits, &targets).unwrap();
1371        assert_array_eq(
1372            loss_mean,
1373            expected_mean,
1374            tolerances::MLX_DEFAULT.rtol,
1375            tolerances::MLX_DEFAULT.atol,
1376        );
1377
1378        let nll_loss = NllLossBuilder::new()
1379            .reduction(LossReduction::Sum)
1380            .build()
1381            .unwrap();
1382        let loss_sum = nll_loss.apply(&logits, &targets).unwrap();
1383        assert_array_eq(
1384            loss_sum,
1385            expected_sum,
1386            tolerances::MLX_DEFAULT.rtol,
1387            tolerances::MLX_DEFAULT.atol,
1388        );
1389    }
1390
1391    #[test]
1392    fn test_gaussian_nll_loss() {
1393        let inputs = array!([[0.1, 0.2], [0.3, 0.4]]);
1394        let targets = array!([[0.2, 0.1], [0.1, 0.2]]);
1395        let vars = array!([[0.1, 0.2], [0.3, 0.4]]);
1396
1397        // Test with reduction 'none', full=False
1398        let gaussian_nll_loss = GaussianNllLossBuilder::new()
1399            .full(false)
1400            .reduction(LossReduction::None)
1401            .build()
1402            .unwrap();
1403        let loss_none = gaussian_nll_loss.apply(&inputs, &targets, &vars).unwrap();
1404        let expected_none = array!([[-1.101293, -0.779719], [-0.535320, -0.408145]]);
1405        assert_array_eq(
1406            loss_none,
1407            &expected_none,
1408            tolerances::MLX_DEFAULT.rtol,
1409            tolerances::MLX_DEFAULT.atol,
1410        );
1411
1412        // Test with reduction 'mean', full=False
1413        let gaussian_nll_loss = GaussianNllLossBuilder::new()
1414            .full(false)
1415            .reduction(LossReduction::Mean)
1416            .build()
1417            .unwrap();
1418        let loss_mean = gaussian_nll_loss.apply(&inputs, &targets, &vars).unwrap();
1419        let expected_mean = expected_none.mean(None).unwrap();
1420        assert_array_eq(
1421            loss_mean,
1422            expected_mean,
1423            tolerances::MLX_DEFAULT.rtol,
1424            tolerances::MLX_DEFAULT.atol,
1425        );
1426
1427        // Test with reduction 'sum', full=False
1428        let gaussian_nll_loss = GaussianNllLossBuilder::new()
1429            .full(false)
1430            .reduction(LossReduction::Sum)
1431            .build()
1432            .unwrap();
1433        let loss_sum = gaussian_nll_loss.apply(&inputs, &targets, &vars).unwrap();
1434        let expected_sum = expected_none.sum(None).unwrap();
1435        assert_array_eq(
1436            loss_sum,
1437            expected_sum,
1438            tolerances::MLX_DEFAULT.rtol,
1439            tolerances::MLX_DEFAULT.atol,
1440        );
1441
1442        // Test with reduction='none', full=True
1443        let gaussian_nll_loss = GaussianNllLossBuilder::new()
1444            .full(true)
1445            .reduction(LossReduction::None)
1446            .build()
1447            .unwrap();
1448        let loss_none_full = gaussian_nll_loss.apply(&inputs, &targets, &vars).unwrap();
1449        let expected_none_full = array!([[-0.182354, 0.139220], [0.383619, 0.510793]]);
1450        assert_array_eq(
1451            loss_none_full,
1452            &expected_none_full,
1453            tolerances::MLX_DEFAULT.rtol,
1454            tolerances::MLX_DEFAULT.atol,
1455        );
1456
1457        // Test with reduction='mean', full=True
1458        let gaussian_nll_loss = GaussianNllLossBuilder::new()
1459            .full(true)
1460            .reduction(LossReduction::Mean)
1461            .build()
1462            .unwrap();
1463        let loss_mean_full = gaussian_nll_loss.apply(&inputs, &targets, &vars).unwrap();
1464        let expected_mean_full = expected_none_full.mean(None).unwrap();
1465        assert_array_eq(
1466            loss_mean_full,
1467            expected_mean_full,
1468            tolerances::MLX_DEFAULT.rtol,
1469            tolerances::MLX_DEFAULT.atol,
1470        );
1471
1472        // Test with reduction='sum', full=True
1473        let gaussian_nll_loss = GaussianNllLossBuilder::new()
1474            .full(true)
1475            .reduction(LossReduction::Sum)
1476            .build()
1477            .unwrap();
1478        let loss_sum_full = gaussian_nll_loss.apply(&inputs, &targets, &vars).unwrap();
1479        let expected_sum_full = expected_none_full.sum(None).unwrap();
1480        assert_array_eq(
1481            loss_sum_full,
1482            expected_sum_full,
1483            tolerances::MLX_DEFAULT.rtol,
1484            tolerances::MLX_DEFAULT.atol,
1485        );
1486    }
1487
1488    #[test]
1489    fn test_kl_div_loss() {
1490        let p_logits = array!([[0.5, 0.5], [0.8, 0.2]]).log().unwrap();
1491        let q_logits = array!([[0.5, 0.5], [0.2, 0.8]]).log().unwrap();
1492
1493        // Test with reduction 'none'
1494        let kl_div_loss = KlDivLossBuilder::new()
1495            .reduction(LossReduction::None)
1496            .build()
1497            .unwrap();
1498        let loss_none = kl_div_loss.apply(&p_logits, &q_logits).unwrap();
1499        let expected_none = array!([0.0, 0.831777]);
1500        assert_array_eq(
1501            loss_none,
1502            &expected_none,
1503            tolerances::MLX_DEFAULT.rtol,
1504            tolerances::MLX_DEFAULT.atol,
1505        );
1506
1507        // Test with reduction 'mean'
1508        let kl_div_loss = KlDivLossBuilder::new()
1509            .reduction(LossReduction::Mean)
1510            .build()
1511            .unwrap();
1512        let loss_mean = kl_div_loss.apply(&p_logits, &q_logits).unwrap();
1513        let expected_mean = expected_none.mean(None).unwrap();
1514        assert_array_eq(
1515            loss_mean,
1516            expected_mean,
1517            tolerances::MLX_DEFAULT.rtol,
1518            tolerances::MLX_DEFAULT.atol,
1519        );
1520
1521        // Test with reduction 'sum'
1522        let kl_div_loss = KlDivLossBuilder::new()
1523            .reduction(LossReduction::Sum)
1524            .build()
1525            .unwrap();
1526        let loss_sum = kl_div_loss.apply(&p_logits, &q_logits).unwrap();
1527        let expected_sum = expected_none.sum(None).unwrap();
1528        assert_array_eq(
1529            loss_sum,
1530            expected_sum,
1531            tolerances::MLX_DEFAULT.rtol,
1532            tolerances::MLX_DEFAULT.atol,
1533        );
1534    }
1535
1536    #[test]
1537    fn test_triplet_loss() {
1538        let anchors = array!([[1, 2, 3], [1, 2, 3]]);
1539        let positives = array!([[4, 5, 6], [0, -1, 2]]);
1540        let negatives = array!([[7, 8, 9], [3, 2, 3]]);
1541
1542        // Test with reduction 'none'
1543        let triplet_loss = TripletLossBuilder::new()
1544            .reduction(LossReduction::None)
1545            .build()
1546            .unwrap();
1547        let loss_none = triplet_loss
1548            .apply(&anchors, &positives, &negatives)
1549            .unwrap();
1550        let expected_none = array!([0.0, 2.31662]);
1551        assert_array_eq(
1552            loss_none,
1553            &expected_none,
1554            tolerances::MLX_DEFAULT.rtol,
1555            tolerances::MLX_DEFAULT.atol,
1556        );
1557
1558        // Test with reduction 'mean'
1559        let triplet_loss = TripletLossBuilder::new()
1560            .reduction(LossReduction::Mean)
1561            .build()
1562            .unwrap();
1563        let loss_mean = triplet_loss
1564            .apply(&anchors, &positives, &negatives)
1565            .unwrap();
1566        let expected_mean = expected_none.mean(None).unwrap();
1567        assert_array_eq(
1568            loss_mean,
1569            expected_mean,
1570            tolerances::MLX_DEFAULT.rtol,
1571            tolerances::MLX_DEFAULT.atol,
1572        );
1573
1574        // Test with reduction 'sum'
1575        let triplet_loss = TripletLossBuilder::new()
1576            .reduction(LossReduction::Sum)
1577            .build()
1578            .unwrap();
1579        let loss_sum = triplet_loss
1580            .apply(&anchors, &positives, &negatives)
1581            .unwrap();
1582        let expected_sum = expected_none.sum(None).unwrap();
1583        assert_array_eq(
1584            loss_sum,
1585            expected_sum,
1586            tolerances::MLX_DEFAULT.rtol,
1587            tolerances::MLX_DEFAULT.atol,
1588        );
1589    }
1590
1591    #[test]
1592    fn test_hinge_loss() {
1593        let inputs = array!([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]);
1594        let targets = array!([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]);
1595        let hinge_loss = HingeLossBuilder::new()
1596            .reduction(LossReduction::Mean)
1597            .build()
1598            .unwrap();
1599        let loss = hinge_loss.apply(&inputs, &targets).unwrap();
1600        assert_eq!(loss.item_exact::<f32>(), 1.0);
1601    }
1602
1603    #[test]
1604    fn test_huber_loss() {
1605        let inputs = array!([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]);
1606        let targets = array!([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]);
1607        let huber_loss = HuberLossBuilder::new()
1608            .reduction(LossReduction::Mean)
1609            .build()
1610            .unwrap();
1611        let loss = huber_loss.apply(&inputs, &targets).unwrap();
1612        assert_eq!(loss.item_exact::<f32>(), 0.5);
1613    }
1614
1615    #[test]
1616    fn test_log_cosh_loss() {
1617        let inputs = array!([[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]);
1618        let targets = array!([[0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0]]);
1619        let log_cosh_loss = LogCoshLossBuilder::new()
1620            .reduction(LossReduction::Mean)
1621            .build()
1622            .unwrap();
1623        let loss = log_cosh_loss.apply(&inputs, &targets).unwrap();
1624        assert_float_eq!(loss.item_exact::<f32>(), 0.433781, abs <= 1e-6);
1625    }
1626
1627    #[test]
1628    fn test_cosine_similarity_loss() {
1629        let embeddings1 = array!([[0.5, 0.5, 0.2, 0.9], [0.1, 0.3, 0.5, 0.5]]);
1630        let embeddings2 = array!([[0.6, 0.4, 0.3, 0.8], [0.2, 0.5, 0.6, 0.4]]);
1631
1632        // Test with reduction 'none'
1633        let cosine_similarity_loss = CosineSimilarityLossBuilder::new()
1634            .reduction(LossReduction::None)
1635            .build()
1636            .unwrap();
1637        let loss_none = cosine_similarity_loss
1638            .apply(&embeddings1, &embeddings2)
1639            .unwrap();
1640        let expected_none = array!([0.985344, 0.961074]);
1641        assert_array_eq(
1642            loss_none,
1643            &expected_none,
1644            tolerances::MLX_DEFAULT.rtol,
1645            tolerances::MLX_DEFAULT.atol,
1646        );
1647
1648        // Test with reduction 'mean'
1649        let cosine_similarity_loss = CosineSimilarityLossBuilder::new()
1650            .reduction(LossReduction::Mean)
1651            .build()
1652            .unwrap();
1653        let loss_mean = cosine_similarity_loss
1654            .apply(&embeddings1, &embeddings2)
1655            .unwrap();
1656        let expected_mean = expected_none.mean(None).unwrap();
1657        assert_array_eq(
1658            loss_mean,
1659            expected_mean,
1660            tolerances::MLX_DEFAULT.rtol,
1661            tolerances::MLX_DEFAULT.atol,
1662        );
1663
1664        // Test with reduction 'sum'
1665        let cosine_similarity_loss = CosineSimilarityLossBuilder::new()
1666            .reduction(LossReduction::Sum)
1667            .build()
1668            .unwrap();
1669        let loss_sum = cosine_similarity_loss
1670            .apply(&embeddings1, &embeddings2)
1671            .unwrap();
1672        let expected_sum = expected_none.sum(None).unwrap();
1673        assert_array_eq(
1674            loss_sum,
1675            expected_sum,
1676            tolerances::MLX_DEFAULT.rtol,
1677            tolerances::MLX_DEFAULT.atol,
1678        );
1679    }
1680
1681    #[test]
1682    fn test_margin_ranking_loss() {
1683        let inputs1 = array!([-0.573409, -0.765166, -0.0638]);
1684        let inputs2 = array!([0.75596, 0.225763, 0.256995]);
1685        let targets = array!([1, 1, -1]);
1686
1687        // Test with no margin
1688        let margin_ranking_loss = MarginRankingLossBuilder::new()
1689            .reduction(LossReduction::None)
1690            .build()
1691            .unwrap();
1692        let loss = margin_ranking_loss
1693            .apply(&inputs1, &inputs2, &targets)
1694            .unwrap();
1695        let expected = array!([1.329369, 0.990929, 0.0]);
1696        assert_array_eq(
1697            loss,
1698            expected,
1699            tolerances::MLX_DEFAULT.rtol,
1700            tolerances::MLX_DEFAULT.atol,
1701        );
1702
1703        // Test with margin
1704        let margin_ranking_loss = MarginRankingLossBuilder::new()
1705            .margin(0.5)
1706            .reduction(LossReduction::None)
1707            .build()
1708            .unwrap();
1709        let loss = margin_ranking_loss
1710            .apply(&inputs1, &inputs2, &targets)
1711            .unwrap();
1712        let expected = array!([1.829369, 1.490929, 0.179205]);
1713        assert_array_eq(
1714            loss,
1715            expected,
1716            tolerances::MLX_DEFAULT.rtol,
1717            tolerances::MLX_DEFAULT.atol,
1718        );
1719    }
1720}
mlx_rs/losses.rs

mlx_rs/
losses.rs
