Fix unresolved mx.array docstring references

[ayaangazali]

Proposed changes

A few docstrings write the array type as mx.array, which Sphinx cannot resolve, so those types render as dead links. Building with -n reports four of them:

docs/src/python/nn/_autosummary_functions/mlx.nn.losses.cosine_similarity_loss.rst:2:
WARNING: py:class reference target not found: mx.array [ref.class]
docs/src/python/nn/_autosummary/mlx.nn.Module.load_weights.rst:2:
WARNING: py:class reference target not found: mx.array [ref.class]

The rest of the codebase writes plain array in docstrings, which resolves. losses.py is a good example of the inconsistency: 29 places write (array) and only cosine_similarity_loss writes (mx.array), in both its Args and its Returns.

Changed to array in:

• nn/losses.py, cosine_similarity_loss (x1, x2, and the return type)
• nn/layers/base.py, Module.load_weights (file_or_weights)
• optimizers.py, init_single and apply_single (parameter, gradient)

After the change a clean -n build reports zero unresolved mx.array references, down from four.

Note the type annotations in the signatures still say mx.array, which is correct Python since that is how the module is imported. This only touches the docstring text that Sphinx tries to turn into a link.

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works
• I have updated the necessary documentation (if needed)

(docstring text only, no behavior change; test_optimizers.py passes and the docs build is clean of these warnings)

---

process note, same as always: freshman contributor and Claude Code helps me hunt, but i found these by building the docs in nitpick mode, then chased the count from four down to zero one file at a time to make sure i had actually got all of them rather than assuming.