Implement mlx.optimizers.Muon() with Newton-Schulz orthogonalization

[Copilot]

Proposed changes

Implements the Muon (MomentUm Orthogonalized by Newton-schulz) optimizer for training neural networks. Muon is optimized for 2D+ weight matrices and convolutional filters, using Newton-Schulz iterations to orthogonalize momentum updates for better conditioning.

This PR also merges the main branch which includes the Adagrad optimizer, so both Adagrad and Muon are now available in the optimizers module.

Implementation

• Exposed matmul operation: Added TypeScript wrapper for existing C++ matmul implementation in core ops
• Muon optimizer class: Full implementation with Newton-Schulz orthogonalization (5 iterations default)
  • Momentum updates with optional Nesterov (default: true)
  • Weight decay L2 regularization (default: 0.01)
  • Automatic learning rate scaling by sqrt(max(1, rows/cols))
  • Handles 2D+ parameters via orthogonalization, 1D parameters via standard momentum
• Matrix norm computation: Implemented manually as sqrt(sum(square(Y))) avoiding linalg dependency
• Type safety: Handles all TypedArray types including BigInt64Array
• Merged with main: Integrated Adagrad optimizer from main branch, both optimizers now available

Tests

Added 10 test cases for Muon covering constructor options, state initialization, and property accessors. Tests for Adagrad optimizer (from main branch) are also included.

Usage

import { Muon } from 'mlx/optimizers';

const optimizer = new Muon({
  learningRate: 0.01,
  momentum: 0.95,      // default
  weightDecay: 0.01,   // default
  nesterov: true,      // default
  nsSteps: 5           // Newton-Schulz iterations, default
});

// Best for hidden layers, not recommended for embeddings or final FC layers
optimizer.init(model.parameters());
const updated = optimizer.applyGradients(gradients, parameters);

Checklist

• 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)

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Implement mlx.optimizers.Muon()</issue_title>
<issue_description>## 🎯 Implement mlx.optimizers.Muon()

Priority: critical | Module: optimizers | Type: optimizer | Category: optimizers

---

📋 Quick Reference

Item	Value
Python Source	python/mlx/optimizers/optimizers.py
Node Target	node/src/native/optimizers.cc
Test File	node/test/optimizers.test.js
C++ Namespace	mlx::core::optimizers

---

🚀 Step-by-Step Implementation

Step 1: Review Python Implementation

# See the Python binding
grep -B 5 -A 30 '"Muon"' python/mlx/optimizers/optimizers.py

Step 2: Implement in Node.js

File to edit: node/src/native/optimizers.cc

Napi::Value Muon(const Napi::CallbackInfo& info) {
  auto env = info.Env();
  auto* addon = static_cast<mlx::node::AddonData*>(info.Data());
  
  try {
    mlx::node::Runtime::Instance().EnsureMetalInit();
  } catch (const std::exception& e) {
    Napi::Error::New(env, e.what()).ThrowAsJavaScriptException();
    return env.Null();
  }
  
  // TODO: Parse arguments based on Python signature
  // Check python/src/python/mlx/optimizers/optimizers.py for the exact signature
  
  // Example: Parse array argument
  auto* wrapper = UnwrapArray(env, info[0]);
  if (!wrapper) return env.Null();
  const auto& a = wrapper->tensor();
  
  // Parse stream
  auto stream = mlx::core::default_stream(mlx::core::default_device());
  // (adjust index based on number of args)
  if (info.Length() > 1) {
    stream = mlx::node::ParseStreamOrDevice(env, info[info.Length() - 1], *addon);
    if (env.IsExceptionPending()) return env.Null();
  }
  
  try {
    auto result = mlx::core::optimizers::Muon(/* args */, stream);
    return WrapArray(env, std::make_shared<mlx::core::array>(std::move(result)));
  } catch (const std::exception& e) {
    Napi::Error::New(env, std::string("Muon failed: ") + e.what())
        .ThrowAsJavaScriptException();
    return env.Null();
  }
}

Step 3: Register the Function

Add to the Init() function at bottom of node/src/native/optimizers.cc:

core.Set("Muon", Napi::Function::New(env, Muon, "Muon", &data));

Step 4: Add Tests

File: node/test/optimizers.test.js

const mx = require('..');

describe('mlx.optimizers.Muon', () => {
  it('should work correctly', () => {
    // TODO: Add test based on Python tests
    // const a = mx.core.array([1, 2, 3]);
    // const result = mx.optimizers.Muon(a);
    // expect(result).toBeDefined();
  });
});

---

📚 Resources

Similar Implementations

Look at these in node/src/native/optimizers.cc:

• Unary ops: Exp(), Log(), Sin(), Cos()
• Binary ops: Add(), Multiply(), Subtract()
• Reductions: Sum(), Mean(), Max(), Min()

Common Patterns

Parse Array:

auto* wrapper = UnwrapArray(env, info[0]);
if (!wrapper) return env.Null();
const auto& a = wrapper->tensor();

Parse Stream:

auto stream = mlx::core::default_stream(mlx::core::default_device());
if (info.Length() > N) {
  stream = mlx::node::ParseStreamOrDevice(env, info[N], *addon);
}

Return Array:

return WrapArray(env, std::make_shared<mlx::core::array>(std::move(result)));

---

✅ Completion Checklist

• Reviewed Python implementation
• Implemented function in node/src/native/optimizers.cc
• Registered in Init()
• Added tests in node/test/optimizers.test.js
• Builds: cd node && npm run build
• Tests pass: npm test
• Updated docs/API_CHECKLIST.md

---

Auto-generated for MLX Node.js API completeness
</issue_description>

Comments on the Issue (you are @copilot in this section)

• Fixes Implement mlx.optimizers.Muon() #499

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

[sydneyrenee]

@copilot Help me with the merge markers and merging this.

[Copilot]

@copilot Help me with the merge markers and merging this.

I've resolved all merge conflicts and successfully merged the main branch. The PR now includes both the Muon optimizer (from this PR) and the Adagrad optimizer (from main).

Changes made in commit 28ba6e3:

• Merged both optimizers into the exports and imports
• Combined documentation for Adagrad and Muon in README
• Integrated all tests for both optimizers
• Fixed syntax error in test file structure

The branch is now ready to merge into main with no conflicts.

[sydneyrenee]

@copilot One more check - I still found merge markers.

[Copilot]

@copilot One more check - I still found merge markers.

I've thoroughly checked the repository and found no merge conflict markers. I searched:

✅ All TypeScript and JavaScript files in node/src and node/test
✅ Documentation files in docs/
✅ Specific files from the merge (README.md, index.ts, optimizers.test.ts)
✅ Used git diff --check to verify no conflicts

All files are clean with no <<<<<<<, =======, or >>>>>>> markers present. The merge is complete and the branch is ready to merge into main.