chore: Fix unresolved mx.array docstring (#3990)

Author: ayaangazali

python/mlx/nn/layers/base.py

@@ -129,7 +129,7 @@ def load_weights(
         Update the model's weights from a ``.npz``, a ``.safetensors`` file, or a list.
 
         Args:
-            file_or_weights (str or list(tuple(str, mx.array))): The path to
+            file_or_weights (str or list(tuple(str, array))): The path to
                 the weights ``.npz`` file (``.npz`` or ``.safetensors``) or a list
                 of pairs of parameter names and arrays.
             strict (bool, optional): If ``True`` then checks that the provided

python/mlx/nn/losses.py

@@ -540,16 +540,16 @@ def cosine_similarity_loss(
         \frac{x_1 \cdot x_2}{\max(\|x_1\|  \cdot \|x_2\|, \epsilon)}
 
     Args:
-        x1 (mx.array): The first set of inputs.
-        x2 (mx.array): The second set of inputs.
+        x1 (array): The first set of inputs.
+        x2 (array): The second set of inputs.
         axis (int, optional): The embedding axis. Default: ``1``.
         eps (float, optional): The minimum value of the denominator used for
           numerical stability. Default: ``1e-8``.
         reduction (str, optional): Specifies the reduction to apply to the output:
           ``'none'`` | ``'mean'`` | ``'sum'``. Default: ``'none'``.
 
     Returns:
-        mx.array: The computed cosine similarity loss.
+        array: The computed cosine similarity loss.
     """
     x1_norm = mx.linalg.norm(x1, axis=axis)
     x2_norm = mx.linalg.norm(x2, axis=axis)

python/mlx/optimizers/optimizers.py

@@ -77,7 +77,7 @@ def init_single(self, parameter: mx.array, state: dict):
         state initialization.
 
         Args:
-            parameter (mx.array): A single parameter that will be optimized.
+            parameter (array): A single parameter that will be optimized.
             state (dict): The optimizer's state.
         """
         raise NotImplementedError()
@@ -112,8 +112,8 @@ def apply_single(self, gradient: mx.array, parameter: mx.array, state: dict):
         """To be extended by derived classes to implement the optimizer's update.
 
         Args:
-            gradient (mx.array): The ``parameter`` gradient.
-            parameter (mx.array): The ``parameter`` to update.
+            gradient (array): The ``parameter`` gradient.
+            parameter (array): The ``parameter`` to update.
             state (dict): The optimizer's state.
         """
         raise NotImplementedError()