mlx_rs/nn/
activation.rs

1use std::f32::consts::PI;
2
3use crate::module::{Module, Param};
4use crate::ops::logsumexp_axis;
5use crate::{
6    array,
7    error::{Exception, Result},
8    ops::{abs, exp, maximum, minimum, multiply, select},
9    transforms::compile::compile,
10    Array, Dtype,
11};
12use mlx_internal_macros::{generate_builder, Buildable, Builder};
13use mlx_macros::ModuleParameters;
14
15fn activation_dtype(x: &Array) -> Dtype {
16    if x.dtype().is_inexact() {
17        x.dtype()
18    } else {
19        Dtype::Float32
20    }
21}
22
23/// Applies the element-wise sigmoid logistic sigmoid.
24///
25/// For details, please see
26/// [this documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.sigmoid.html)
27pub fn sigmoid(x: impl AsRef<Array>) -> Result<Array> {
28    crate::ops::sigmoid(x.as_ref())
29}
30
31/// Applies the Rectified Linear Unit.
32///
33/// This is:
34///
35/// ```text
36/// maximum(x, 0)
37/// ```
38pub fn relu(x: impl AsRef<Array>) -> Result<Array> {
39    crate::ops::maximum(x.as_ref(), &array!(0).as_dtype(x.as_ref().dtype())?)
40}
41
42/// Applies the Leaky Rectified Linear Unit.
43///
44/// `neg_slope` is default to 0.01 if not provided.
45///
46/// This is:
47///
48/// ```text
49/// maximum(neg_slope * x, x)
50/// ```
51pub fn leaky_relu(x: impl AsRef<Array>, neg_slope: impl Into<Option<f32>>) -> Result<Array> {
52    let x = x.as_ref();
53    let dtype = activation_dtype(x);
54    let neg_slope = array!(neg_slope.into().unwrap_or(0.01)).as_dtype(dtype)?;
55    // We have to use this indirection, otherwise the compiler cannot
56    // infer the lifetime of the value returned by the closure properly
57    compiled_leaky_relu(x, &neg_slope)
58}
59
60/// Applies the Log Softmax function.
61///
62/// This is:
63///
64/// ```text
65/// x - logsumexp_axis(x, axis, true)
66/// ```
67pub fn log_softmax(x: impl AsRef<Array>, axis: impl Into<Option<i32>>) -> Result<Array> {
68    let x = x.as_ref();
69    let axis = axis.into().unwrap_or(-1);
70    x.subtract(logsumexp_axis(x, axis, true)?)
71}
72
73/// Applies the Exponential Linear Unit.
74///
75/// This is:
76///
77/// ```text
78/// select(x.gt(0), x, alpha * (exp(x) - 1))
79/// ```
80///
81/// # Params
82///
83/// - `x`: The input array
84/// - `alpha`: Default to 1.0 if not provided
85pub fn elu(x: impl AsRef<Array>, alpha: impl Into<Option<f32>>) -> Result<Array> {
86    let x = x.as_ref();
87    let dtype = activation_dtype(x);
88    let alpha = array!(alpha.into().unwrap_or(1.0)).as_dtype(dtype)?;
89    // We have to use this indirection, otherwise the compiler cannot
90    // infer the lifetime of the value returned by the closure properly
91    compiled_elu(x, &alpha)
92}
93
94/// Applies the Rectified Linear Unit 6.
95///
96/// This is:
97///
98/// ```text
99/// minimum(maximum(x, 0), 6)
100/// ```
101pub fn relu6(x: impl AsRef<Array>) -> Result<Array> {
102    compiled_relu6(x.as_ref())
103}
104
105/// Applies the Exponential Linear Unit.
106///
107/// This is:
108///
109/// ```text
110/// logaddexp(x, 0)
111/// ```
112pub fn softplus(x: impl AsRef<Array>) -> Result<Array> {
113    crate::ops::logaddexp(x.as_ref(), &array!(0).as_dtype(x.as_ref().dtype())?)
114}
115
116/// Applies the Softsign function.
117///
118/// This is:
119///
120/// ```text
121/// x / (1 + abs(x))
122/// ```
123pub fn softsign(x: impl AsRef<Array>) -> Result<Array> {
124    compiled_softsign(x.as_ref())
125}
126
127/// Applies the Continuously Differentiable Exponential Linear Unit.
128///
129/// This is:
130///
131/// ```text
132/// maximum(x, 0) + alpha * (exp(minimum(x, 0) / alpha) - 1)
133/// ```
134pub fn celu(x: impl AsRef<Array>, alpha: impl Into<Option<f32>>) -> Result<Array> {
135    let x = x.as_ref();
136    let dtype = activation_dtype(x);
137    let alpha = array!(alpha.into().unwrap_or(1.0)).as_dtype(dtype)?;
138    // We have to use this indirection, otherwise the compiler cannot
139    // infer the lifetime of the value returned by the closure properly
140    compiled_celu(x, &alpha)
141}
142
143/// Applies the Sigmoid Linear Unit. Also known as Swish.
144///
145/// This is:
146///
147/// ```text
148/// x * sigmoid(x)
149/// ```
150pub fn silu(x: impl AsRef<Array>) -> Result<Array> {
151    compiled_silu(x.as_ref())
152}
153
154/// Applies the Log Sigmoid function.
155///
156/// This is:
157///
158/// ```text
159/// -softplus(-x)
160/// ```
161pub fn log_sigmoid(x: impl AsRef<Array>) -> Result<Array> {
162    compiled_log_sigmoid(x.as_ref())
163}
164
165/// Applies the Gaussian Error Linear Units function.
166///
167/// This is:
168///
169/// ```text
170/// x * (1 + erf(x / 2.sqrt())) / 2
171/// ```
172pub fn gelu(x: impl AsRef<Array>) -> Result<Array> {
173    compiled_gelu(x.as_ref())
174}
175
176/// An approximation to Gaussian Error Linear Unit.
177///
178/// This is:
179///
180/// ```text
181/// 0.5 * x * (1 + tanh(sqrt(2 / PI) * (x + 0.044715 * x ** 3)))
182/// ```
183pub fn gelu_approximate(x: impl AsRef<Array>) -> Result<Array> {
184    compiled_gelu_approximate(x.as_ref())
185}
186
187/// A fast approximation to Gaussian Error Linear Unit.
188///
189/// This is:
190///
191/// ```text
192/// x * sigmoid(1.702 * x)
193/// ```
194pub fn gelu_fast_approximate(x: impl AsRef<Array>) -> Result<Array> {
195    compiled_gelu_fast_approximate(x.as_ref())
196}
197
198/// Applies the gated linear unit function.
199///
200/// This function splits the `axis` dimension of the input into two halves
201/// (`a` and `b`) and applies `a * sigmoid(b)`.
202pub fn glu(x: impl AsRef<Array>, axis: impl Into<Option<i32>>) -> Result<Array> {
203    let split = x.as_ref().split_equal(2, axis)?;
204    let (a, b) = (&split[0], &split[1]);
205    Ok(a * sigmoid(b)?)
206}
207
208/// Applies the Step Activation Function.
209///
210/// This function implements a binary step activation, where the output is set
211/// to 1 if the input is greater than or equal to a specified threshold, and 0 otherwise.
212///
213/// This is:
214///
215/// ```text
216/// select(x.ge(threshold), 1, 0)
217/// ```
218pub fn step(x: impl AsRef<Array>, threshold: impl Into<Option<f32>>) -> Result<Array> {
219    let x = x.as_ref();
220    let dtype = activation_dtype(x);
221    let threshold = array!(threshold.into().unwrap_or(0.0)).as_dtype(dtype)?;
222    crate::ops::select(&x.ge(threshold)?, &array!(1), &array!(0))
223}
224
225/// Applies the Scaled Exponential Linear Unit.
226///
227/// This is:
228///
229/// ```text
230/// elu(x, 1.67326) * 1.0507
231/// ```
232pub fn selu(x: impl AsRef<Array>) -> Result<Array> {
233    compiled_selu(x.as_ref())
234}
235
236/// Applies the element-wise parametric ReLU.
237///
238/// This is:
239///
240/// ```text
241/// maximum(0, x) + alpha * minimum(0, x)
242/// ```
243pub fn prelu(x: impl AsRef<Array>, alpha: impl AsRef<Array>) -> Result<Array> {
244    compiled_prelu(x.as_ref(), alpha.as_ref())
245}
246
247/// Applies the Mish function, element-wise.
248///
249/// Mish: A Self Regularized Non-Monotonic Neural Activation Function.
250///
251/// Reference: [https://arxiv.org/abs/1908.08681](https://arxiv.org/abs/1908.08681)
252///
253/// This is:
254///
255/// ```text
256/// x * tanh(softplus(x))
257/// ```
258pub fn mish(x: impl AsRef<Array>) -> Result<Array> {
259    compiled_mish(x.as_ref())
260}
261
262/// Applies the hardswish function, element-wise.
263///
264/// This is:
265///
266/// ```text
267/// x * minimum(maximum(x + 3, 0), 6) / 6
268/// ```
269pub fn hard_swish(x: impl AsRef<Array>) -> Result<Array> {
270    compiled_hard_swish(x.as_ref())
271}
272
273generate_builder! {
274    /// Applies the gated linear unit function.
275    ///
276    /// This splits the `axis` dimension of the input into two halves
277    /// (`a` and `b`) and applies `a * sigmoid(b)`.
278    #[derive(Debug, Clone, ModuleParameters, Buildable)]
279    #[module(root = crate)]
280    #[buildable(root = crate)]
281    #[builder(root = crate)]
282    pub struct Glu {
283        /// The axis to split the input tensor. Default to [`Glu::DEFAULT_AXIS`] if not provided.
284        #[builder(optional, default = Glu::DEFAULT_AXIS)]
285        pub axis: i32,
286    }
287}
288
289impl Glu {
290    /// The default axis value.
291    pub const DEFAULT_AXIS: i32 = -1;
292}
293
294impl Module<&Array> for Glu {
295    type Error = Exception;
296    type Output = Array;
297
298    fn forward(&mut self, x: &Array) -> Result<Array> {
299        glu(x, self.axis)
300    }
301
302    fn training_mode(&mut self, _: bool) {}
303}
304
305/// Applies the element-wise logistic sigmoid.
306///
307/// For details, please see
308/// [this documentation](https://ml-explore.github.io/mlx/build/html/python/_autosummary/mlx.core.sigmoid.html)
309#[derive(Debug, Clone, ModuleParameters)]
310#[module(root = crate)]
311pub struct Sigmoid;
312
313impl Module<&Array> for Sigmoid {
314    type Error = Exception;
315    type Output = Array;
316
317    fn forward(&mut self, x: &Array) -> Result<Array> {
318        sigmoid(x)
319    }
320
321    fn training_mode(&mut self, _: bool) {}
322}
323
324/// Applies the Mish function, element-wise.
325///
326/// Mish: A Self Regularized Non-Monotonic Neural Activation Function.
327///
328/// Reference: [https://arxiv.org/abs/1908.08681](https://arxiv.org/abs/1908.08681)
329///
330/// This is:
331///
332/// ```text
333/// x * tanh(softplus(x))
334/// ```
335#[derive(Debug, Clone, ModuleParameters)]
336#[module(root = crate)]
337pub struct Mish;
338
339impl Module<&Array> for Mish {
340    type Error = Exception;
341    type Output = Array;
342
343    fn forward(&mut self, x: &Array) -> Result<Array> {
344        mish(x)
345    }
346
347    fn training_mode(&mut self, _: bool) {}
348}
349
350/// Applies the Rectified Linear Unit.
351///
352/// This is:
353///
354/// ```text
355/// maximum(x, 0)
356/// ```
357#[derive(Debug, Clone, ModuleParameters)]
358#[module(root = crate)]
359pub struct Relu;
360
361impl Module<&Array> for Relu {
362    type Error = Exception;
363    type Output = Array;
364
365    fn forward(&mut self, x: &Array) -> Result<Array> {
366        relu(x)
367    }
368
369    fn training_mode(&mut self, _: bool) {}
370}
371
372generate_builder! {
373    /// Applies the Leaky Rectified Linear Unit.
374    ///
375    /// This is:
376    ///
377    /// ```text
378    /// maximum(neg_slope * x, x)
379    /// ```
380    #[derive(Debug, Clone, ModuleParameters, Buildable)]
381    #[module(root = crate)]
382    #[buildable(root = crate)]
383    #[builder(root = crate)]
384    pub struct LeakyRelu {
385        /// The negative slope. Default to [`LeakyReLU::DEFAULT_NEG_SLOPE`] if not provided.
386        #[builder(optional, default = LeakyRelu::DEFAULT_NEG_SLOPE)]
387        pub neg_slope: f32,
388    }
389}
390
391impl LeakyRelu {
392    /// The default negative slope value.
393    pub const DEFAULT_NEG_SLOPE: f32 = 0.01;
394}
395
396impl Module<&Array> for LeakyRelu {
397    type Error = Exception;
398    type Output = Array;
399
400    fn forward(&mut self, x: &Array) -> Result<Array> {
401        leaky_relu(x, self.neg_slope)
402    }
403
404    fn training_mode(&mut self, _: bool) {}
405}
406
407/// Applies the Rectified Linear Unit 6.
408///
409/// This is:
410///
411/// ```text
412/// minimum(maximum(x, 0), 6)
413/// ```
414#[derive(Debug, Clone, ModuleParameters)]
415#[module(root = crate)]
416pub struct Relu6;
417
418impl Module<&Array> for Relu6 {
419    type Error = Exception;
420    type Output = Array;
421
422    fn forward(&mut self, x: &Array) -> Result<Array> {
423        relu6(x)
424    }
425
426    fn training_mode(&mut self, _: bool) {}
427}
428
429generate_builder! {
430    /// Applies the Softmax function.
431    ///
432    /// This is:
433    ///
434    /// ```text
435    /// exp(x - logsumexp(x, axis, true))
436    /// ```
437    #[derive(Debug, Clone, ModuleParameters, Buildable)]
438    #[module(root = crate)]
439    #[buildable(root = crate)]
440    #[builder(root = crate)]
441    pub struct Softmax {
442        /// The axis to apply the softmax.
443        #[builder(optional, default = Softmax::DEFAULT_AXIS)]
444        pub axis: i32,
445    }
446}
447
448impl Softmax {
449    /// The default axis value.
450    pub const DEFAULT_AXIS: i32 = -1;
451}
452
453impl Module<&Array> for Softmax {
454    type Error = Exception;
455    type Output = Array;
456
457    fn forward(&mut self, x: &Array) -> Result<Array> {
458        crate::ops::softmax_axis(x, self.axis, None)
459    }
460
461    fn training_mode(&mut self, _: bool) {}
462}
463
464/// Applies the Softplus function.
465///
466/// This is:
467///
468/// ```text
469/// logaddexp(x, 0)
470/// ```
471#[derive(Debug, Clone, ModuleParameters)]
472#[module(root = crate)]
473pub struct Softplus;
474
475impl Module<&Array> for Softplus {
476    type Error = Exception;
477    type Output = Array;
478
479    fn forward(&mut self, x: &Array) -> Result<Array> {
480        softplus(x)
481    }
482
483    fn training_mode(&mut self, _: bool) {}
484}
485
486/// Applies the Softsign function.
487///
488/// This is:
489///
490/// ```text
491/// x / (1 + abs(x))
492/// ```
493#[derive(Debug, Clone, ModuleParameters)]
494#[module(root = crate)]
495pub struct Softsign;
496
497impl Module<&Array> for Softsign {
498    type Error = Exception;
499    type Output = Array;
500
501    fn forward(&mut self, x: &Array) -> Result<Array> {
502        softsign(x)
503    }
504
505    fn training_mode(&mut self, _: bool) {}
506}
507
508generate_builder! {
509    /// Applies the Continuously Differentiable Exponential Linear Unit.
510    ///
511    /// This is:
512    ///
513    /// ```text
514    /// maximum(x, 0) + alpha * (exp(minimum(x, 0) / alpha) - 1)
515    /// ```
516    #[derive(Debug, Clone, ModuleParameters, Buildable)]
517    #[module(root = crate)]
518    #[buildable(root = crate)]
519    #[builder(root = crate)]
520    pub struct Celu {
521        /// The alpha value. Default to [`Celu::DEFAULT_ALPHA`] if not provided.
522        #[builder(optional, default = Celu::DEFAULT_ALPHA)]
523        pub alpha: f32,
524    }
525}
526
527impl Celu {
528    /// The default alpha value.
529    pub const DEFAULT_ALPHA: f32 = 1.0;
530}
531
532impl Module<&Array> for Celu {
533    type Error = Exception;
534    type Output = Array;
535
536    fn forward(&mut self, x: &Array) -> Result<Array> {
537        celu(x, self.alpha)
538    }
539
540    fn training_mode(&mut self, _: bool) {}
541}
542
543/// Applies the Sigmoid Linear Unit. Also known as Swish.
544///
545/// This is:
546///
547/// ```text
548/// x * sigmoid(x)
549/// ```
550#[derive(Debug, Clone, ModuleParameters)]
551#[module(root = crate)]
552pub struct Silu;
553
554impl Module<&Array> for Silu {
555    type Error = Exception;
556    type Output = Array;
557
558    fn forward(&mut self, x: &Array) -> Result<Array> {
559        silu(x)
560    }
561
562    fn training_mode(&mut self, _: bool) {}
563}
564
565generate_builder! {
566    /// Applies the Log Softmax function.
567    ///
568    /// This is:
569    ///
570    /// ```text
571    /// x - logsumexp(x, axis, true)
572    /// ```
573    #[derive(Debug, Clone, ModuleParameters, Buildable)]
574    #[module(root = crate)]
575    #[buildable(root = crate)]
576    #[builder(root = crate)]
577    pub struct LogSoftmax {
578        /// The axis value. Default to [`LogSoftmax::DEFAULT_AXIS`] if not provided.
579        #[builder(optional, default = LogSoftmax::DEFAULT_AXIS)]
580        pub axis: i32,
581    }
582}
583
584impl LogSoftmax {
585    /// The default axis value.
586    pub const DEFAULT_AXIS: i32 = -1;
587}
588
589impl Module<&Array> for LogSoftmax {
590    type Error = Exception;
591    type Output = Array;
592
593    fn forward(&mut self, x: &Array) -> Result<Array> {
594        log_softmax(x, self.axis)
595    }
596
597    fn training_mode(&mut self, _: bool) {}
598}
599
600/// Applies the Log Sigmoid function.
601///
602/// This is:
603///
604/// ```text
605/// -softplus(-x)
606/// ```
607#[derive(Debug, Clone, ModuleParameters)]
608#[module(root = crate)]
609pub struct LogSigmoid;
610
611impl Module<&Array> for LogSigmoid {
612    type Error = Exception;
613    type Output = Array;
614
615    fn forward(&mut self, x: &Array) -> Result<Array> {
616        log_sigmoid(x)
617    }
618
619    fn training_mode(&mut self, _: bool) {}
620}
621
622/// Applies the element-wise parametric ReLU.
623///
624/// This is:
625///
626/// ```text
627/// maximum(0, x) + alpha * minimum(0, x)
628/// ```
629#[derive(Debug, Clone, ModuleParameters, Buildable)]
630#[module(root = crate)]
631#[buildable(root = crate)]
632pub struct Prelu {
633    /// The alpha value. See [`prelu`] for more details.
634    #[param]
635    #[builder(ignore)]
636    pub weight: Param<Array>, // TODO: double check if this is trainable
637}
638
639/// The builder for the Prelu module.
640#[derive(Debug, Clone, Builder)]
641#[builder(
642    root = crate,
643    build_with = build_prelu,
644    default_infallible,
645    err = Exception,
646)]
647pub struct PreluBuilder {
648    /// The count. Default to [`Prelu::DEFAULT_COUNT`] if not provided.
649    #[builder(optional, default = Prelu::DEFAULT_COUNT)]
650    pub count: i32,
651
652    /// The value. Default to [`Prelu::DEFAULT_VALUE`] if not provided.
653    #[builder(optional, default = Prelu::DEFAULT_VALUE)]
654    pub value: f32,
655}
656
657/// Builds the Prelu module.
658fn build_prelu(builder: PreluBuilder) -> Result<Prelu> {
659    let count = builder.count;
660    let value = builder.value;
661    let weight = Param::new(crate::ops::full::<f32>(&[count], &array!(value))?);
662    Ok(Prelu { weight })
663}
664
665impl Prelu {
666    /// The default count value.
667    pub const DEFAULT_COUNT: i32 = 1;
668
669    /// The default value.
670    pub const DEFAULT_VALUE: f32 = 0.25;
671}
672
673impl Module<&Array> for Prelu {
674    type Error = Exception;
675    type Output = Array;
676
677    fn forward(&mut self, x: &Array) -> Result<Array> {
678        prelu(x, &self.weight)
679    }
680
681    fn training_mode(&mut self, _: bool) {}
682}
683
684/// Variants of Gaussian Error Linear Units function.
685#[derive(Debug, Clone, Copy, Default)]
686pub enum GeluApprox {
687    /// Uses [`gelu`]
688    #[default]
689    None,
690
691    /// Uses [`gelu_approximate`]
692    Precise,
693
694    /// Uses [`gelu_fast_approximate`]
695    Fast,
696}
697
698generate_builder! {
699    /// Applies the Gaussian Error Linear Units function.
700    ///
701    /// There are three variants:
702    ///
703    /// - `GeluApprox::None`: Uses [`gelu`]. This is the default.
704    /// - `GeluApprox::Precise`: Uses [`gelu_approximate`]
705    /// - `GeluApprox::Fast`: Uses [`gelu_fast_approximate`]
706    #[derive(Debug, Clone, ModuleParameters, Buildable)]
707    #[module(root = crate)]
708    #[buildable(root = crate)]
709    #[builder(root = crate)]
710    pub struct Gelu {
711        /// The approximation to use. Default to `GeluApprox::None` if not provided.
712        #[builder(optional, default = GeluApprox::None)]
713        pub approximate: GeluApprox,
714    }
715}
716
717impl Module<&Array> for Gelu {
718    type Error = Exception;
719    type Output = Array;
720
721    fn forward(&mut self, x: &Array) -> Result<Array> {
722        match self.approximate {
723            GeluApprox::None => gelu(x),
724            GeluApprox::Precise => gelu_approximate(x),
725            GeluApprox::Fast => gelu_fast_approximate(x),
726        }
727    }
728
729    fn training_mode(&mut self, _: bool) {}
730}
731
732/// Applies the hyperbolic tangent function
733#[derive(Debug, Clone, ModuleParameters)]
734#[module(root = crate)]
735pub struct Tanh;
736
737impl Module<&Array> for Tanh {
738    type Error = Exception;
739    type Output = Array;
740
741    fn forward(&mut self, x: &Array) -> Result<Array> {
742        crate::ops::tanh(x)
743    }
744
745    fn training_mode(&mut self, _: bool) {}
746}
747
748/// Applies the hardswish function, element-wise
749///
750/// This is:
751///
752/// ```text
753/// x * minimum(maximum(x + 3, 0), 6) / 6
754/// ```
755#[derive(Debug, Clone, ModuleParameters)]
756#[module(root = crate)]
757pub struct HardSwish;
758
759impl Module<&Array> for HardSwish {
760    type Error = Exception;
761    type Output = Array;
762
763    fn forward(&mut self, x: &Array) -> Result<Array> {
764        hard_swish(x)
765    }
766
767    fn training_mode(&mut self, _: bool) {}
768}
769
770generate_builder! {
771    /// Applies the Step Activation Function.
772    ///
773    /// This function implements a binary step activation, where the output is set
774    /// to 1 if the input is greater than or equal to a specified threshold, and 0 otherwise.
775    ///
776    /// This is:
777    ///
778    /// ```text
779    /// select(x.ge(threshold), 1, 0)
780    /// ```
781    #[derive(Debug, Clone, ModuleParameters, Buildable)]
782    #[module(root = crate)]
783    #[buildable(root = crate)]
784    #[builder(root = crate)]
785    pub struct Step {
786        /// The threshold value. Default to [`Step::DEFAULT_THRESHOLD`] if not provided.
787        #[builder(optional, default = Step::DEFAULT_THRESHOLD)]
788        pub threshold: f32,
789    }
790}
791
792impl Step {
793    /// The default threshold value.
794    pub const DEFAULT_THRESHOLD: f32 = 0.0;
795}
796
797impl Module<&Array> for Step {
798    type Error = Exception;
799    type Output = Array;
800
801    fn forward(&mut self, x: &Array) -> Result<Array> {
802        step(x, self.threshold)
803    }
804
805    fn training_mode(&mut self, _: bool) {}
806}
807
808/// Applies the Scaled Exponential Linear Unit.
809///
810/// This is:
811///
812/// ```text
813/// elu(x, 1.67326) * 1.0507
814/// ```
815#[derive(Debug, Clone, ModuleParameters)]
816#[module(root = crate)]
817pub struct Selu;
818
819impl Module<&Array> for Selu {
820    type Error = Exception;
821    type Output = Array;
822
823    fn forward(&mut self, x: &Array) -> Result<Array> {
824        selu(x)
825    }
826
827    fn training_mode(&mut self, _: bool) {}
828}
829
830/* -------------------------------------------------------------------------- */
831/*                        Compiled activation functions                       */
832/* -------------------------------------------------------------------------- */
833
834#[inline]
835fn compiled_leaky_relu(x: &Array, neg_slope: &Array) -> Result<Array> {
836    let f = |(x_, neg_slope_): (&Array, &Array)| {
837        // This will not panic because a scalar can always be broadcasted to any shape
838        let a = multiply(neg_slope_, x_)?;
839        maximum(&a, x_)
840    };
841    let mut compiled = compile(f, true);
842    compiled((x, neg_slope))
843}
844
845#[inline]
846fn compiled_elu(x: &Array, alpha: &Array) -> Result<Array> {
847    let f = |(x_, alpha_): (&Array, &Array)| {
848        select(
849            &x_.gt(&array!(0).as_dtype(x_.dtype())?)?,
850            x_,
851            alpha_ * (exp(x_)? - array!(1).as_dtype(alpha_.dtype())?),
852        )
853    };
854    let mut compiled = compile(f, true);
855    compiled((x, alpha))
856}
857
858#[inline]
859fn compiled_relu6(x: &Array) -> Result<Array> {
860    let f = |x_: &Array| {
861        let dtype = activation_dtype(x_);
862        minimum(
863            maximum(x_, &array!(0).as_dtype(x_.dtype())?)?,
864            &array!(6.0).as_dtype(dtype)?,
865        )
866    };
867    let mut compiled = compile(f, true);
868    compiled(x)
869}
870
871#[inline]
872fn compiled_softsign(x: &Array) -> Result<Array> {
873    let f = |x_: &Array| x_.divide(array!(1).as_dtype(x_.dtype())? + abs(x_)?);
874    let mut compiled = compile(f, true);
875    compiled(x)
876}
877
878#[inline]
879fn compiled_celu(x: &Array, alpha: &Array) -> Result<Array> {
880    let f = |(x_, alpha_): (&Array, &Array)| {
881        maximum(x_, &array!(0.0).as_dtype(alpha_.dtype())?)?.add(alpha_.multiply(
882            exp(&(minimum(x_, &array!(0.0).as_dtype(alpha_.dtype())?)? / alpha_))?
883                - array!(1).as_dtype(alpha_.dtype())?,
884        )?)
885    };
886    let mut compiled = compile(f, true);
887    compiled((x, alpha))
888}
889
890#[inline]
891fn compiled_silu(x: &Array) -> Result<Array> {
892    let f = |x_: &Array| x_.multiply(sigmoid(x_)?);
893    let mut compiled = compile(f, true);
894    compiled(x)
895}
896
897#[inline]
898fn compiled_log_sigmoid(x: &Array) -> Result<Array> {
899    let f = |x_: &Array| Ok(-softplus(&(-x_))?);
900    let mut compiled = compile(f, true);
901    compiled(x)
902}
903
904#[inline]
905fn compiled_gelu(x: &Array) -> Result<Array> {
906    use crate::ops::erf;
907    let f = |x_: &Array| {
908        let dtype = activation_dtype(x_);
909        x_.multiply(
910            array!(1).as_dtype(dtype)? + erf(&(x_ / array!(2f32.sqrt()).as_dtype(dtype)?))?,
911        )?
912        .divide(array!(2).as_dtype(dtype)?)
913    };
914    let mut compiled = compile(f, true);
915    compiled(x)
916}
917
918#[inline]
919fn compiled_gelu_approximate(x: &Array) -> Result<Array> {
920    use crate::ops::tanh;
921
922    let f = move |x_: &Array| {
923        let dtype = activation_dtype(x_);
924        array!(0.5).as_dtype(dtype)?.multiply(x_)?.multiply(
925            array!(1).as_dtype(dtype)?.add(tanh(
926                &(array!((2.0 / PI).sqrt()).as_dtype(dtype)?.multiply(
927                    x_ + array!(0.044715)
928                        .as_dtype(dtype)?
929                        .multiply(x_.power(&array!(3).as_dtype(x_.dtype())?)?)?,
930                )?),
931            )?)?,
932        )
933    };
934    let mut compiled = compile(f, true);
935    compiled(x)
936}
937
938#[inline]
939fn compiled_gelu_fast_approximate(x: &Array) -> Result<Array> {
940    let f = |x_: &Array| {
941        let dtype = activation_dtype(x_);
942        x_.multiply(sigmoid(&(array!(1.702).as_dtype(dtype)? * x_))?)
943    };
944    let mut compiled = compile(f, true);
945    compiled(x)
946}
947
948#[inline]
949fn compiled_selu(x: &Array) -> Result<Array> {
950    let f = |x_: &Array| {
951        let y = elu(x_, 1.67326)?;
952        y.multiply(array!(1.0507).as_dtype(y.dtype())?)
953    };
954    let mut compiled = compile(f, true);
955    compiled(x)
956}
957
958#[inline]
959fn compiled_prelu(x: &Array, alpha: &Array) -> Result<Array> {
960    let f = |(x_, alpha_): (&Array, &Array)| {
961        maximum(&array!(0).as_dtype(x_.dtype())?, x_)?
962            .add(alpha_ * minimum(&array!(0).as_dtype(x_.dtype())?, x_)?)
963    };
964    let mut compiled = compile(f, true);
965    compiled((x, alpha))
966}
967
968#[inline]
969fn compiled_mish(x: &Array) -> Result<Array> {
970    use crate::ops::tanh;
971
972    let f = |x_: &Array| x_.multiply(tanh(&softplus(x_)?)?);
973    let mut compiled = compile(f, true);
974    compiled(x)
975}
976
977#[inline]
978fn compiled_hard_swish(x: &Array) -> Result<Array> {
979    let f = |x_: &Array| {
980        let max_x_plus_3 = maximum(
981            &(x_ + array!(3).as_dtype(x_.dtype())?),
982            &array!(0).as_dtype(x_.dtype())?,
983        )?;
984        x_.multiply(minimum(&max_x_plus_3, &array!(6).as_dtype(x_.dtype())?)?)?
985            .divide(&array!(6).as_dtype(x_.dtype())?)
986    };
987    let mut compiled = compile(f, true);
988    compiled(x)
989}
990
991#[cfg(test)]
992mod tests {
993    use crate::{builder::Builder, random::uniform, Dtype};
994    use float_eq::assert_float_eq;
995
996    use super::*;
997
998    #[test]
999    fn test_activation_dtypes_python() {
1000        use crate::{with_device, Device};
1001        use Dtype::{Bfloat16, Float16, Float32, Int32};
1002
1003        // Python MLX 0.32.2, CPU: nn.<function>(mx.array([-3, -1, 0, 1, 7], dtype=dtype)).
1004        // Rows are bfloat16, float16, float32, int32; prelu uses mx.array(0.25, dtype=dtype).
1005        // Python calls the GELU approximations gelu_approx/gelu_fast_approx and hard_swish hardswish.
1006        type Activation = fn(&Array) -> Result<Array>;
1007        type Case = (&'static str, Activation, [Dtype; 4], [[f32; 5]; 4]);
1008        let floating = [Bfloat16, Float16, Float32, Float32];
1009        let cases: &[Case] = &[
1010            (
1011                "relu",
1012                |x| relu(x),
1013                [Bfloat16, Float16, Float32, Int32],
1014                [
1015                    [0.0, 0.0, 0.0, 1.0, 7.0],
1016                    [0.0, 0.0, 0.0, 1.0, 7.0],
1017                    [0.0, 0.0, 0.0, 1.0, 7.0],
1018                    [0.0, 0.0, 0.0, 1.0, 7.0],
1019                ],
1020            ),
1021            (
1022                "leaky_relu",
1023                |x| leaky_relu(x, None),
1024                floating,
1025                [
1026                    [-0.030029297, -0.010009766, 0.0, 1.0, 7.0],
1027                    [-0.02999878, -0.010002136, 0.0, 1.0, 7.0],
1028                    [-0.03, -0.01, 0.0, 1.0, 7.0],
1029                    [-0.03, -0.01, 0.0, 1.0, 7.0],
1030                ],
1031            ),
1032            (
1033                "elu",
1034                |x| elu(x, None),
1035                floating,
1036                [
1037                    [-0.94921875, -0.6328125, 0.0, 1.0, 7.0],
1038                    [-0.9501953, -0.63183594, 0.0, 1.0, 7.0],
1039                    [-0.95021296, -0.63212055, 0.0, 1.0, 7.0],
1040                    [-0.95021296, -0.63212055, 0.0, 1.0, 7.0],
1041                ],
1042            ),
1043            (
1044                "relu6",
1045                |x| relu6(x),
1046                floating,
1047                [
1048                    [0.0, 0.0, 0.0, 1.0, 6.0],
1049                    [0.0, 0.0, 0.0, 1.0, 6.0],
1050                    [0.0, 0.0, 0.0, 1.0, 6.0],
1051                    [0.0, 0.0, 0.0, 1.0, 6.0],
1052                ],
1053            ),
1054            (
1055                "softplus",
1056                |x| softplus(x),
1057                floating,
1058                [
1059                    [0.048583984, 0.3125, 0.69140625, 1.3125, 7.0],
1060                    [0.048583984, 0.31323242, 0.6933594, 1.3134766, 7.0],
1061                    [0.048587356, 0.3132617, 0.6931472, 1.3132617, 7.000911],
1062                    [0.048587356, 0.3132617, 0.6931472, 1.3132617, 7.000911],
1063                ],
1064            ),
1065            (
1066                "softsign",
1067                |x| softsign(x),
1068                floating,
1069                [
1070                    [-0.75, -0.5, 0.0, 0.5, 0.875],
1071                    [-0.75, -0.5, 0.0, 0.5, 0.875],
1072                    [-0.75, -0.5, 0.0, 0.5, 0.875],
1073                    [-0.75, -0.5, 0.0, 0.5, 0.875],
1074                ],
1075            ),
1076            (
1077                "celu",
1078                |x| celu(x, None),
1079                floating,
1080                [
1081                    [-0.94921875, -0.6328125, 0.0, 1.0, 7.0],
1082                    [-0.9501953, -0.63183594, 0.0, 1.0, 7.0],
1083                    [-0.95021296, -0.63212055, 0.0, 1.0, 7.0],
1084                    [-0.95021296, -0.63212055, 0.0, 1.0, 7.0],
1085                ],
1086            ),
1087            (
1088                "silu",
1089                |x| silu(x),
1090                floating,
1091                [
1092                    [-0.14257812, -0.26953125, 0.0, 0.73046875, 7.0],
1093                    [-0.14233398, -0.26879883, 0.0, 0.73095703, 6.9921875],
1094                    [-0.14227763, -0.26894143, 0.0, 0.7310586, 6.993623],
1095                    [-0.14227763, -0.26894143, 0.0, 0.7310586, 6.993623],
1096                ],
1097            ),
1098            (
1099                "log_sigmoid",
1100                |x| log_sigmoid(x),
1101                floating,
1102                [
1103                    [-3.046875, -1.3125, -0.69140625, -0.3125, -0.00091171265],
1104                    [
1105                        -3.0488281,
1106                        -1.3134766,
1107                        -0.6933594,
1108                        -0.31323242,
1109                        -0.0009112358,
1110                    ],
1111                    [
1112                        -3.0485873,
1113                        -1.3132617,
1114                        -0.6931472,
1115                        -0.3132617,
1116                        -0.0009114663,
1117                    ],
1118                    [
1119                        -3.0485873,
1120                        -1.3132617,
1121                        -0.6931472,
1122                        -0.3132617,
1123                        -0.0009114663,
1124                    ],
1125                ],
1126            ),
1127            (
1128                "gelu",
1129                |x| gelu(x),
1130                floating,
1131                [
1132                    [-0.005859375, -0.15820312, 0.0, 0.84375, 7.0],
1133                    [-0.0043945312, -0.1586914, 0.0, 0.8413086, 7.0],
1134                    [-0.004049957, -0.15865538, 0.0, 0.8413446, 7.0],
1135                    [-0.004049957, -0.15865538, 0.0, 0.8413446, 7.0],
1136                ],
1137            ),
1138            (
1139                "gelu_approximate",
1140                |x| gelu_approximate(x),
1141                floating,
1142                [
1143                    [-0.005859375, -0.15820312, 0.0, 0.84375, 7.0],
1144                    [-0.0036621094, -0.15893555, 0.0, 0.8408203, 7.0],
1145                    [-0.003637433, -0.158808, 0.0, 0.841192, 7.0],
1146                    [-0.003637433, -0.158808, 0.0, 0.841192, 7.0],
1147                ],
1148            ),
1149            (
1150                "gelu_fast_approximate",
1151                |x| gelu_fast_approximate(x),
1152                floating,
1153                [
1154                    [-0.017822266, -0.15429688, 0.0, 0.84765625, 7.0],
1155                    [-0.018081665, -0.1541748, 0.0, 0.8457031, 7.0],
1156                    [-0.018071312, -0.15420423, 0.0, 0.84579575, 6.9999533],
1157                    [-0.018071312, -0.15420423, 0.0, 0.84579575, 6.9999533],
1158                ],
1159            ),
1160            (
1161                "step",
1162                |x| step(x, None),
1163                [Int32; 4],
1164                [
1165                    [0.0, 0.0, 1.0, 1.0, 1.0],
1166                    [0.0, 0.0, 1.0, 1.0, 1.0],
1167                    [0.0, 0.0, 1.0, 1.0, 1.0],
1168                    [0.0, 0.0, 1.0, 1.0, 1.0],
1169                ],
1170            ),
1171            (
1172                "selu",
1173                |x| selu(x),
1174                floating,
1175                [
1176                    [-1.6640625, -1.1015625, 0.0, 1.046875, 7.3125],
1177                    [-1.6708984, -1.1103516, 0.0, 1.0507812, 7.3554688],
1178                    [-1.6705638, -1.1113275, 0.0, 1.0507, 7.3548994],
1179                    [-1.6705638, -1.1113275, 0.0, 1.0507, 7.3548994],
1180                ],
1181            ),
1182            (
1183                "prelu",
1184                |x| prelu(x, array!(0.25).as_dtype(x.dtype())?),
1185                [Bfloat16, Float16, Float32, Int32],
1186                [
1187                    [-0.75, -0.25, 0.0, 1.0, 7.0],
1188                    [-0.75, -0.25, 0.0, 1.0, 7.0],
1189                    [-0.75, -0.25, 0.0, 1.0, 7.0],
1190                    [0.0, 0.0, 0.0, 1.0, 7.0],
1191                ],
1192            ),
1193            (
1194                "mish",
1195                |x| mish(x),
1196                floating,
1197                [
1198                    [-0.14550781, -0.30273438, 0.0, 0.86328125, 7.0],
1199                    [-0.14562988, -0.3034668, 0.0, 0.8652344, 7.0],
1200                    [-0.14564747, -0.30340147, 0.0, 0.8650984, 6.9999886],
1201                    [-0.14564747, -0.30340147, 0.0, 0.8650984, 6.9999886],
1202                ],
1203            ),
1204            (
1205                "hard_swish",
1206                |x| hard_swish(x),
1207                floating,
1208                [
1209                    [-0.0, -0.33398438, 0.0, 0.66796875, 7.0],
1210                    [-0.0, -0.33325195, 0.0, 0.6665039, 7.0],
1211                    [-0.0, -0.33333334, 0.0, 0.6666667, 7.0],
1212                    [0.0, -0.33333334, 0.0, 0.6666667, 7.0],
1213                ],
1214            ),
1215        ];
1216        with_device(Device::cpu(), || {
1217            for &(name, activation, output_dtypes, expected) in cases {
1218                for ((dtype, expected_dtype), values) in [Bfloat16, Float16, Float32, Int32]
1219                    .into_iter()
1220                    .zip(output_dtypes)
1221                    .zip(expected)
1222                {
1223                    let x = array!([-3, -1, 0, 1, 7]).as_dtype(dtype).unwrap();
1224                    let y = activation(&x).unwrap();
1225                    assert_eq!(y.dtype(), expected_dtype, "{name}({dtype:?})");
1226                    #[cfg(feature = "metal")]
1227                    with_device(Device::gpu(), || {
1228                        let y = activation(&x).unwrap();
1229                        assert_eq!(y.dtype(), expected_dtype, "{name}({dtype:?}) on Metal");
1230                        y.eval().unwrap();
1231                    });
1232                    assert_eq!(y.shape(), &[5]);
1233                    for (got, want) in y.to_vec_cast::<f32>().unwrap().into_iter().zip(values) {
1234                        assert_float_eq!(got, want, abs <= 1e-5 * (1.0 + want.abs()));
1235                    }
1236                }
1237            }
1238        });
1239    }
1240
1241    #[test]
1242    fn test_step_fractional_threshold_python() {
1243        // Python MLX 0.32.2: nn.step(mx.array([-1, 0, 0.1, 1], dtype=dtype), 0.1).
1244        for dtype in [
1245            Dtype::Bfloat16,
1246            Dtype::Float16,
1247            Dtype::Float32,
1248            Dtype::Int32,
1249        ] {
1250            let x = array!([-1.0, 0.0, 0.1, 1.0]).as_dtype(dtype).unwrap();
1251            let y = step(&x, 0.1).unwrap();
1252            assert_eq!(y.dtype(), Dtype::Int32);
1253            let expected = if dtype == Dtype::Int32 {
1254                [0, 0, 0, 1]
1255            } else {
1256                [0, 0, 1, 1]
1257            };
1258            assert_eq!(y.to_vec_exact::<i32>().unwrap(), expected);
1259        }
1260    }
1261
1262    #[test]
1263    fn test_prelu_float32_weights_python() {
1264        // Python MLX 0.32.2: nn.prelu(mx.array([-3, -1, 0, 1, 7], dtype=dtype), mx.array(0.25)).
1265        for dtype in [
1266            Dtype::Bfloat16,
1267            Dtype::Float16,
1268            Dtype::Float32,
1269            Dtype::Int32,
1270        ] {
1271            let x = array!([-3, -1, 0, 1, 7]).as_dtype(dtype).unwrap();
1272            let y = prelu(&x, array!(0.25)).unwrap();
1273            assert_eq!(y.dtype(), Dtype::Float32);
1274            assert_eq!(
1275                y.to_vec_exact::<f32>().unwrap(),
1276                [-0.75, -0.25, 0.0, 1.0, 7.0]
1277            );
1278        }
1279    }
1280
1281    // The following tests are ported from the swift binding:
1282    // mlx-swift/Tests/MLXTests/IntegrationTests.swift
1283    #[test]
1284    fn test_glu() {
1285        crate::random::seed(850).unwrap();
1286        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1287        assert_eq!(a.shape(), &[2, 8, 16]);
1288        assert_eq!(a.dtype(), Dtype::Float32);
1289        assert_float_eq!(
1290            a.mean(None).unwrap().item_exact::<f32>(),
1291            0.547_252_66,
1292            abs <= 0.010_945_053
1293        );
1294        assert_float_eq!(
1295            a.sum(None).unwrap().item_exact::<f32>(),
1296            140.096_68,
1297            abs <= 2.801_933_5
1298        );
1299        let result = Glu::new().forward(&a).unwrap();
1300        assert_eq!(result.shape(), &[2, 8, 8]);
1301        assert_eq!(result.dtype(), Dtype::Float32);
1302        assert_float_eq!(
1303            result.mean(None).unwrap().item_exact::<f32>(),
1304            0.333_276_75,
1305            abs <= 0.006_665_535
1306        );
1307        assert_float_eq!(
1308            result.sum(None).unwrap().item_exact::<f32>(),
1309            42.659_424,
1310            abs <= 0.853_188_46
1311        );
1312    }
1313
1314    #[test]
1315    fn test_sigmoid() {
1316        crate::random::seed(589).unwrap();
1317        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1318        assert_eq!(a.shape(), &[2, 8, 16]);
1319        assert_eq!(a.dtype(), Dtype::Float32);
1320        assert_float_eq!(
1321            a.mean(None).unwrap().item_exact::<f32>(),
1322            0.529_697_9,
1323            abs <= 0.010_593_958
1324        );
1325        assert_float_eq!(
1326            a.sum(None).unwrap().item_exact::<f32>(),
1327            135.602_66,
1328            abs <= 2.712_053_3
1329        );
1330        let result = Sigmoid.forward(&a).unwrap();
1331        assert_eq!(result.shape(), &[2, 8, 16]);
1332        assert_eq!(result.dtype(), Dtype::Float32);
1333        assert_float_eq!(
1334            result.mean(None).unwrap().item_exact::<f32>(),
1335            0.627_014,
1336            abs <= 0.012_540_28
1337        );
1338        assert_float_eq!(
1339            result.sum(None).unwrap().item_exact::<f32>(),
1340            160.515_58,
1341            abs <= 3.210_311_7
1342        );
1343    }
1344
1345    #[test]
1346    fn test_mish() {
1347        crate::random::seed(122).unwrap();
1348        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1349        assert_eq!(a.shape(), &[2, 8, 16]);
1350        assert_eq!(a.dtype(), Dtype::Float32);
1351        assert_float_eq!(
1352            a.mean(None).unwrap().item_exact::<f32>(),
1353            0.501_719_8,
1354            abs <= 0.010_034_395
1355        );
1356        assert_float_eq!(
1357            a.sum(None).unwrap().item_exact::<f32>(),
1358            128.440_26,
1359            abs <= 2.568_805_2
1360        );
1361        let result = Mish.forward(&a).unwrap();
1362        assert_eq!(result.shape(), &[2, 8, 16]);
1363        assert_eq!(result.dtype(), Dtype::Float32);
1364        assert_float_eq!(
1365            result.mean(None).unwrap().item_exact::<f32>(),
1366            0.395_375_73,
1367            abs <= 0.007_907_514
1368        );
1369        assert_float_eq!(
1370            result.sum(None).unwrap().item_exact::<f32>(),
1371            101.216_19,
1372            abs <= 2.024_323_7
1373        );
1374    }
1375
1376    #[test]
1377    fn test_relu() {
1378        crate::random::seed(400).unwrap();
1379        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1380        assert_eq!(a.shape(), &[2, 8, 16]);
1381        assert_eq!(a.dtype(), Dtype::Float32);
1382        assert_float_eq!(
1383            a.mean(None).unwrap().item_exact::<f32>(),
1384            0.478_322_74,
1385            abs <= 0.009_566_455
1386        );
1387        assert_float_eq!(
1388            a.sum(None).unwrap().item_exact::<f32>(),
1389            122.450_62,
1390            abs <= 2.449_012_5
1391        );
1392        let result = Relu.forward(&a).unwrap();
1393        assert_eq!(result.shape(), &[2, 8, 16]);
1394        assert_eq!(result.dtype(), Dtype::Float32);
1395        assert_float_eq!(
1396            result.mean(None).unwrap().item_exact::<f32>(),
1397            0.478_322_74,
1398            abs <= 0.009_566_455
1399        );
1400        assert_float_eq!(
1401            result.sum(None).unwrap().item_exact::<f32>(),
1402            122.450_62,
1403            abs <= 2.449_012_5
1404        );
1405    }
1406
1407    #[test]
1408    fn test_leaky_relu() {
1409        crate::random::seed(93).unwrap();
1410        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1411        assert_eq!(a.shape(), &[2, 8, 16]);
1412        assert_eq!(a.dtype(), Dtype::Float32);
1413        assert_float_eq!(
1414            a.mean(None).unwrap().item_exact::<f32>(),
1415            0.499_930_68,
1416            abs <= 0.009_998_614
1417        );
1418        assert_float_eq!(
1419            a.sum(None).unwrap().item_exact::<f32>(),
1420            127.982_254,
1421            abs <= 2.559_645_2
1422        );
1423        let result = LeakyRelu::new().forward(&a).unwrap();
1424        assert_eq!(result.shape(), &[2, 8, 16]);
1425        assert_eq!(result.dtype(), Dtype::Float32);
1426        assert_float_eq!(
1427            result.mean(None).unwrap().item_exact::<f32>(),
1428            0.499_930_68,
1429            abs <= 0.009_998_614
1430        );
1431        assert_float_eq!(
1432            result.sum(None).unwrap().item_exact::<f32>(),
1433            127.982_254,
1434            abs <= 2.559_645_2
1435        );
1436    }
1437
1438    #[test]
1439    fn test_relu6() {
1440        crate::random::seed(379).unwrap();
1441        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1442        assert_eq!(a.shape(), &[2, 8, 16]);
1443        assert_eq!(a.dtype(), Dtype::Float32);
1444        assert_float_eq!(
1445            a.mean(None).unwrap().item_exact::<f32>(),
1446            0.493_258_66,
1447            abs <= 0.009_865_173
1448        );
1449        assert_float_eq!(
1450            a.sum(None).unwrap().item_exact::<f32>(),
1451            126.274_216,
1452            abs <= 2.525_484_3
1453        );
1454        let result = Relu6.forward(&a).unwrap();
1455        assert_eq!(result.shape(), &[2, 8, 16]);
1456        assert_eq!(result.dtype(), Dtype::Float32);
1457        assert_float_eq!(
1458            result.mean(None).unwrap().item_exact::<f32>(),
1459            0.493_258_66,
1460            abs <= 0.009_865_173
1461        );
1462        assert_float_eq!(
1463            result.sum(None).unwrap().item_exact::<f32>(),
1464            126.274_216,
1465            abs <= 2.525_484_3
1466        );
1467    }
1468
1469    #[test]
1470    fn test_softmax() {
1471        crate::random::seed(853).unwrap();
1472        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1473        assert_eq!(a.shape(), &[2, 8, 16]);
1474        assert_eq!(a.dtype(), Dtype::Float32);
1475        assert_float_eq!(
1476            a.mean(None).unwrap().item_exact::<f32>(),
1477            0.514_396_3,
1478            abs <= 0.010_287_926_5
1479        );
1480        assert_float_eq!(
1481            a.sum(None).unwrap().item_exact::<f32>(),
1482            131.685_46,
1483            abs <= 2.633_709_2
1484        );
1485        let result = Softmax::new().forward(&a).unwrap();
1486        assert_eq!(result.shape(), &[2, 8, 16]);
1487        assert_eq!(result.dtype(), Dtype::Float32);
1488        assert_float_eq!(
1489            result.mean(None).unwrap().item_exact::<f32>(),
1490            0.062_499_996,
1491            abs <= 0.001_25
1492        );
1493        assert_float_eq!(
1494            result.sum(None).unwrap().item_exact::<f32>(),
1495            15.999_999,
1496            abs <= 0.32
1497        );
1498    }
1499
1500    #[test]
1501    fn test_softplus() {
1502        crate::random::seed(118).unwrap();
1503        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1504        assert_eq!(a.shape(), &[2, 8, 16]);
1505        assert_eq!(a.dtype(), Dtype::Float32);
1506        assert_float_eq!(
1507            a.mean(None).unwrap().item_exact::<f32>(),
1508            0.498_981_42,
1509            abs <= 0.009_979_628
1510        );
1511        assert_float_eq!(
1512            a.sum(None).unwrap().item_exact::<f32>(),
1513            127.739_24,
1514            abs <= 2.554_784_8
1515        );
1516        let result = Softplus.forward(&a).unwrap();
1517        assert_eq!(result.shape(), &[2, 8, 16]);
1518        assert_eq!(result.dtype(), Dtype::Float32);
1519        assert_float_eq!(
1520            result.mean(None).unwrap().item_exact::<f32>(),
1521            0.982_857_76,
1522            abs <= 0.019_657_155
1523        );
1524        assert_float_eq!(
1525            result.sum(None).unwrap().item_exact::<f32>(),
1526            251.611_59,
1527            abs <= 5.032_232
1528        );
1529    }
1530
1531    #[test]
1532    fn test_softsign() {
1533        crate::random::seed(37).unwrap();
1534        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1535        assert_eq!(a.shape(), &[2, 8, 16]);
1536        assert_eq!(a.dtype(), Dtype::Float32);
1537        assert_float_eq!(
1538            a.mean(None).unwrap().item_exact::<f32>(),
1539            0.506_551_27,
1540            abs <= 0.010_131_026
1541        );
1542        assert_float_eq!(
1543            a.sum(None).unwrap().item_exact::<f32>(),
1544            129.677_12,
1545            abs <= 2.593_542_6
1546        );
1547        let result = Softsign.forward(&a).unwrap();
1548        assert_eq!(result.shape(), &[2, 8, 16]);
1549        assert_eq!(result.dtype(), Dtype::Float32);
1550        assert_float_eq!(
1551            result.mean(None).unwrap().item_exact::<f32>(),
1552            0.314_089_83,
1553            abs <= 0.006_281_797
1554        );
1555        assert_float_eq!(
1556            result.sum(None).unwrap().item_exact::<f32>(),
1557            80.407,
1558            abs <= 1.608_14
1559        );
1560    }
1561
1562    // The unit test below is adapted from the python binding:
1563    // mlx/python/tests/test_nn.py
1564    #[test]
1565    fn test_celu() {
1566        let x = array!([1.0, -1.0, 0.0]);
1567        let y = Celu::new().forward(&x).unwrap();
1568        let epsilon = array!(1e-4);
1569        let expected_y = array!([1.0, -0.6321, 0.0]);
1570        assert!(y
1571            .subtract(&expected_y)
1572            .unwrap()
1573            .abs()
1574            .unwrap()
1575            .lt(&epsilon)
1576            .unwrap()
1577            .all(None)
1578            .unwrap()
1579            .item_exact::<bool>());
1580        assert_eq!(y.shape(), &[3]);
1581        assert_eq!(y.dtype(), Dtype::Float32);
1582
1583        let y = CeluBuilder::new()
1584            .alpha(1.1_f32)
1585            .build()
1586            .unwrap()
1587            .forward(&x)
1588            .unwrap();
1589        let expected_y = array!([1.0, -0.6568, 0.0]);
1590        assert!(y
1591            .subtract(&expected_y)
1592            .unwrap()
1593            .abs()
1594            .unwrap()
1595            .lt(&epsilon)
1596            .unwrap()
1597            .all(None)
1598            .unwrap()
1599            .item_exact::<bool>());
1600        assert_eq!(y.shape(), &[3]);
1601        assert_eq!(y.dtype(), Dtype::Float32);
1602    }
1603
1604    #[test]
1605    fn test_silu() {
1606        crate::random::seed(22).unwrap();
1607        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1608        assert_eq!(a.shape(), &[2, 8, 16]);
1609        assert_eq!(a.dtype(), Dtype::Float32);
1610        assert_float_eq!(
1611            a.mean(None).unwrap().item_exact::<f32>(),
1612            0.502_970_6,
1613            abs <= 0.010_059_412
1614        );
1615        assert_float_eq!(
1616            a.sum(None).unwrap().item_exact::<f32>(),
1617            128.760_47,
1618            abs <= 2.575_209_4
1619        );
1620        let result = Silu.forward(&a).unwrap();
1621        assert_eq!(result.shape(), &[2, 8, 16]);
1622        assert_eq!(result.dtype(), Dtype::Float32);
1623        assert_float_eq!(
1624            result.mean(None).unwrap().item_exact::<f32>(),
1625            0.331_970_93,
1626            abs <= 0.006_639_418_7
1627        );
1628        assert_float_eq!(
1629            result.sum(None).unwrap().item_exact::<f32>(),
1630            84.984_56,
1631            abs <= 1.699_691_2
1632        );
1633    }
1634
1635    #[test]
1636    fn test_log_softmax() {
1637        crate::random::seed(199).unwrap();
1638        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1639        assert_eq!(a.shape(), &[2, 8, 16]);
1640        assert_eq!(a.dtype(), Dtype::Float32);
1641        assert_float_eq!(
1642            a.mean(None).unwrap().item_exact::<f32>(),
1643            0.527_843_7,
1644            abs <= 0.010_556_874
1645        );
1646        assert_float_eq!(
1647            a.sum(None).unwrap().item_exact::<f32>(),
1648            135.127_99,
1649            abs <= 2.702_559_7
1650        );
1651        let result = LogSoftmax::new().forward(&a).unwrap();
1652        assert_eq!(result.shape(), &[2, 8, 16]);
1653        assert_eq!(result.dtype(), Dtype::Float32);
1654        assert_float_eq!(
1655            result.mean(None).unwrap().item_exact::<f32>(),
1656            -2.810_954_6,
1657            abs <= 0.056_219_09
1658        );
1659        assert_float_eq!(
1660            result.sum(None).unwrap().item_exact::<f32>(),
1661            -719.604_4,
1662            abs <= 14.392_087
1663        );
1664    }
1665
1666    #[test]
1667    fn test_log_sigmoid() {
1668        crate::random::seed(984).unwrap();
1669        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1670        assert_eq!(a.shape(), &[2, 8, 16]);
1671        assert_eq!(a.dtype(), Dtype::Float32);
1672        assert_float_eq!(
1673            a.mean(None).unwrap().item_exact::<f32>(),
1674            0.510_977_7,
1675            abs <= 0.010_219_553_5
1676        );
1677        assert_float_eq!(
1678            a.sum(None).unwrap().item_exact::<f32>(),
1679            130.810_29,
1680            abs <= 2.616_205_7
1681        );
1682        let result = LogSigmoid.forward(&a).unwrap();
1683        assert_eq!(result.shape(), &[2, 8, 16]);
1684        assert_eq!(result.dtype(), Dtype::Float32);
1685        assert_float_eq!(
1686            result.mean(None).unwrap().item_exact::<f32>(),
1687            -0.479_598_55,
1688            abs <= 0.009_591_971
1689        );
1690        assert_float_eq!(
1691            result.sum(None).unwrap().item_exact::<f32>(),
1692            -122.777_23,
1693            abs <= 2.455_544_5
1694        );
1695    }
1696
1697    #[test]
1698    fn test_prelu() {
1699        crate::random::seed(993).unwrap();
1700        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1701        assert_eq!(a.shape(), &[2, 8, 16]);
1702        assert_eq!(a.dtype(), Dtype::Float32);
1703        assert_float_eq!(
1704            a.mean(None).unwrap().item_exact::<f32>(),
1705            0.496_651_44,
1706            abs <= 0.009_933_028
1707        );
1708        assert_float_eq!(
1709            a.sum(None).unwrap().item_exact::<f32>(),
1710            127.142_77,
1711            abs <= 2.542_855_3
1712        );
1713        let result = Prelu::new().forward(&a).unwrap();
1714        assert_eq!(result.shape(), &[2, 8, 16]);
1715        assert_eq!(result.dtype(), Dtype::Float32);
1716        assert_float_eq!(
1717            result.mean(None).unwrap().item_exact::<f32>(),
1718            0.496_651_44,
1719            abs <= 0.009_933_028
1720        );
1721        assert_float_eq!(
1722            result.sum(None).unwrap().item_exact::<f32>(),
1723            127.142_77,
1724            abs <= 2.542_855_3
1725        );
1726    }
1727
1728    #[test]
1729    fn test_gelu() {
1730        crate::random::seed(189).unwrap();
1731        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1732        assert_eq!(a.shape(), &[2, 8, 16]);
1733        assert_eq!(a.dtype(), Dtype::Float32);
1734        assert_float_eq!(
1735            a.mean(None).unwrap().item_exact::<f32>(),
1736            0.492_950_32,
1737            abs <= 0.009_859_007
1738        );
1739        assert_float_eq!(
1740            a.sum(None).unwrap().item_exact::<f32>(),
1741            126.195_28,
1742            abs <= 2.523_905_8
1743        );
1744        let result = Gelu::new().forward(&a).unwrap();
1745        assert_eq!(result.shape(), &[2, 8, 16]);
1746        assert_eq!(result.dtype(), Dtype::Float32);
1747        assert_float_eq!(
1748            result.mean(None).unwrap().item_exact::<f32>(),
1749            0.365_638_38,
1750            abs <= 0.007_312_767_7
1751        );
1752        assert_float_eq!(
1753            result.sum(None).unwrap().item_exact::<f32>(),
1754            93.603_424,
1755            abs <= 1.872_068_5
1756        );
1757    }
1758
1759    #[test]
1760    fn test_tanh() {
1761        crate::random::seed(735).unwrap();
1762        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1763        assert_eq!(a.shape(), &[2, 8, 16]);
1764        assert_eq!(a.dtype(), Dtype::Float32);
1765        assert_float_eq!(
1766            a.mean(None).unwrap().item_exact::<f32>(),
1767            0.474_122_7,
1768            abs <= 0.009_482_454_5
1769        );
1770        assert_float_eq!(
1771            a.sum(None).unwrap().item_exact::<f32>(),
1772            121.375_41,
1773            abs <= 2.427_508_4
1774        );
1775        let result = Tanh.forward(&a).unwrap();
1776        assert_eq!(result.shape(), &[2, 8, 16]);
1777        assert_eq!(result.dtype(), Dtype::Float32);
1778        assert_float_eq!(
1779            result.mean(None).unwrap().item_exact::<f32>(),
1780            0.413_079_68,
1781            abs <= 0.008_261_594
1782        );
1783        assert_float_eq!(
1784            result.sum(None).unwrap().item_exact::<f32>(),
1785            105.748_4,
1786            abs <= 2.114_968
1787        );
1788    }
1789
1790    #[test]
1791    fn test_hardswish() {
1792        crate::random::seed(126).unwrap();
1793        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1794        assert_eq!(a.shape(), &[2, 8, 16]);
1795        assert_eq!(a.dtype(), Dtype::Float32);
1796        assert_float_eq!(
1797            a.mean(None).unwrap().item_exact::<f32>(),
1798            0.491_892_46,
1799            abs <= 0.009_837_849
1800        );
1801        assert_float_eq!(
1802            a.sum(None).unwrap().item_exact::<f32>(),
1803            125.924_47,
1804            abs <= 2.518_489_4
1805        );
1806        let result = HardSwish.forward(&a).unwrap();
1807        assert_eq!(result.shape(), &[2, 8, 16]);
1808        assert_eq!(result.dtype(), Dtype::Float32);
1809        assert_float_eq!(
1810            result.mean(None).unwrap().item_exact::<f32>(),
1811            0.299_602_24,
1812            abs <= 0.005_992_044_7
1813        );
1814        assert_float_eq!(
1815            result.sum(None).unwrap().item_exact::<f32>(),
1816            76.698_17,
1817            abs <= 1.533_963_4
1818        );
1819    }
1820
1821    #[test]
1822    fn test_step() {
1823        crate::random::seed(490).unwrap();
1824        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1825        assert_eq!(a.shape(), &[2, 8, 16]);
1826        assert_eq!(a.dtype(), Dtype::Float32);
1827        assert_float_eq!(
1828            a.mean(None).unwrap().item_exact::<f32>(),
1829            0.479_360_64,
1830            abs <= 0.009_587_212_5
1831        );
1832        assert_float_eq!(
1833            a.sum(None).unwrap().item_exact::<f32>(),
1834            122.716_324,
1835            abs <= 2.454_326_4
1836        );
1837        let result = Step::new().forward(&a).unwrap();
1838        assert_eq!(result.shape(), &[2, 8, 16]);
1839        assert_eq!(result.dtype(), Dtype::Int32);
1840        assert_float_eq!(
1841            result.mean(None).unwrap().item_exact::<f32>(),
1842            1.0,
1843            abs <= 0.02
1844        );
1845        assert_eq!(result.sum(None).unwrap().item_exact::<i32>(), 256);
1846    }
1847
1848    #[test]
1849    fn test_selu() {
1850        crate::random::seed(215).unwrap();
1851        let a = uniform::<_, f32>(0.0, 1.0, &[2, 8, 16], None).unwrap();
1852        assert_eq!(a.shape(), &[2, 8, 16]);
1853        assert_eq!(a.dtype(), Dtype::Float32);
1854        assert_float_eq!(
1855            a.mean(None).unwrap().item_exact::<f32>(),
1856            0.493_026_8,
1857            abs <= 0.009_860_536
1858        );
1859        assert_float_eq!(
1860            a.sum(None).unwrap().item_exact::<f32>(),
1861            126.214_86,
1862            abs <= 2.524_297_2
1863        );
1864        let result = Selu.forward(&a).unwrap();
1865        assert_eq!(result.shape(), &[2, 8, 16]);
1866        assert_eq!(result.dtype(), Dtype::Float32);
1867        assert_float_eq!(
1868            result.mean(None).unwrap().item_exact::<f32>(),
1869            0.518_023_2,
1870            abs <= 0.010_360_463_5
1871        );
1872        assert_float_eq!(
1873            result.sum(None).unwrap().item_exact::<f32>(),
1874            132.613_94,
1875            abs <= 2.652_278_7
1876        );
1877    }
1878}
mlx_rs/nn/activation.rs

mlx_rs/nn/
activation.rs
