Add ReLUSquared Activation Function (#250)

* feat: Add ReLUSquared activation function and refactor relu implementation

* feat: Document ReLUSquared activation function in activations and MLXNN files

Author: johnmai-dev

Source/MLXNN/Activations.swift

@@ -33,7 +33,7 @@ public func sigmoid(_ x: MLXArray) -> MLXArray {
 /// - <doc:activations>
 /// - ``ReLU``
 public func relu(_ x: MLXArray) -> MLXArray {
-    maximum(x, 0)
+    compiledRelu(x)
 }
 
 /// Applies the Leaky Rectified Linear Unit.
@@ -103,6 +103,21 @@ public func relu6(_ x: MLXArray) -> MLXArray {
     compiledRelu6(x)
 }
 
+/// Applies the squared Rectified Linear Unit.
+///
+/// This is:
+///
+/// ```swift
+/// MLX.relu(x).square()
+/// ```
+///
+/// /// ### See Also
+/// - <doc:activations>
+/// - ``ReLUSquared``
+public func reluSquared(_ x: MLXArray) -> MLXArray {
+    compiledReluSquared(x)
+}
+
 @available(*, deprecated, renamed: "softplus(_:)")
 @_documentation(visibility: internal)
 public func softPlus(_ x: MLXArray) -> MLXArray {
@@ -458,6 +473,23 @@ open class ReLU6: Module, UnaryLayer {
     }
 }
 
+/// Applies the squared Rectified Linear Unit.
+///
+/// This is:
+///
+/// ```swift
+/// MLX.maximum(x, 0).square()
+///
+///
+/// ### See Also
+/// - <doc:activations>
+/// - ``reluSquared(_:)``
+open class ReLUSquared: Module, UnaryLayer {
+    open func callAsFunction(_ x: MLXArray) -> MLXArray {
+        reluSquared(x)
+    }
+}
+
 @available(*, deprecated, renamed: "Softmax")
 @_documentation(visibility: internal)
 open class SoftMax: Module, UnaryLayer {
@@ -852,3 +884,15 @@ private let compiledHardSwish: @Sendable (MLXArray) -> MLXArray = {
         return x * minimum(maxXPlus3, 6) / 6
     }
 }()
+
+private let compiledRelu: @Sendable (MLXArray) -> MLXArray = {
+    compile(shapeless: true) { x in
+        maximum(x, 0)
+    }
+}()
+
+private let compiledReluSquared: @Sendable (MLXArray) -> MLXArray = {
+    compile(shapeless: true) { x in
+        return relu(x).square()
+    }
+}()

Source/MLXNN/Documentation.docc/MLXNN.md

@@ -128,6 +128,7 @@ These can be used with ``Sequential``.
 - ``prelu(_:alpha:)``
 - ``relu(_:)``
 - ``relu6(_:)``
+- ``reluSquared(_:)``
 - ``selu(_:)``
 - ``silu(_:)``
 - ``sigmoid(_:)``
@@ -148,6 +149,7 @@ These can be used with ``Sequential``.
 - ``PReLU``
 - ``ReLU``
 - ``ReLU6``
+- ``ReLUSquared``
 - ``SELU``
 - ``SiLU``
 - ``Sigmoid``

Source/MLXNN/Documentation.docc/activations.md

@@ -25,6 +25,7 @@ such as `alpha`.
 - ``prelu(_:alpha:)``
 - ``relu(_:)``
 - ``relu6(_:)``
+- ``reluSquared(_:)``
 - ``selu(_:)``
 - ``silu(_:)``
 - ``sigmoid(_:)``
@@ -45,6 +46,7 @@ such as `alpha`.
 - ``PReLU``
 - ``ReLU``
 - ``ReLU6``
+- ``ReLUSquared``
 - ``SELU``
 - ``SiLU``
 - ``Sigmoid``