Implement mlx.optimizers.Adamax()

[Copilot]

Successfully Rebased onto Latest Main

Completed Rebase:

• Rebased copilot/implement-adamax-optimizer onto remotes/origin/main (b09705c)
• Resolved merge conflicts in import statements
  • node/src/optimizers/index.ts: merged to include abs, maximum (Adamax) + matmul, reshape, transpose (Muon)
  • node/test/optimizers.test.ts: merged to include all optimizers (SGD, Adam, Adamax, Lion, Adagrad, RMSprop, Muon)
• Enhanced Adamax documentation to clarify advantages over Adam
• All 6 commits successfully rebased on top of latest main
• Branch properly integrated with Muon Implement mlx.optimizers.Muon() with Newton-Schulz orthogonalization #572 and Adagrad Implement mlx.optimizers.Adagrad() #575

Ready for Review:

• Branch is cleanly rebased and ready to be pushed
• All Adamax functionality preserved and properly integrated
• No merge conflict markers remain

Original prompt

---

This section details on the original issue you should resolve

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

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 '"Adamax"' python/mlx/optimizers/optimizers.py

Step 2: Implement in Node.js

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

Napi::Value Adamax(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::Adamax(/* 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("Adamax 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("Adamax", Napi::Function::New(env, Adamax, "Adamax", &data));

Step 4: Add Tests

File: node/test/optimizers.test.js

const mx = require('..');

describe('mlx.optimizers.Adamax', () => {
  it('should work correctly', () => {
    // TODO: Add test based on Python tests
    // const a = mx.core.array([1, 2, 3]);
    // const result = mx.optimizers.Adamax(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.Adamax() #496

---

✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

[Copilot]

Pull request overview

This PR implements the Adamax optimizer, a variant of Adam that uses an infinity norm instead of L2 norm for the second moment estimation. The implementation extends the Adam class and provides a fully functional optimizer that differs from Adam in its exponentially weighted infinity norm computation.

Key changes:

• Added Adamax optimizer class extending Adam with infinity norm implementation (max(β₂ * v, |g|) vs Adam's β₂ * v + (1-β₂) * g²)
• Implemented comprehensive test suite covering construction, validation, state management, and inheritance behavior
• Updated documentation to reflect the completion of Adamax implementation

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
node/src/optimizers/index.ts	Added Adamax class implementation with constructor, state initialization, and gradient application logic; imported required abs and maximum operations
node/test/optimizers.test.ts	Added comprehensive test suite covering Adamax construction, parameter validation, state initialization, and inheritance verification
node/CHECKLIST.md	Marked Adamax optimizer as fully implemented
docs/OPTIMIZER_API_VERIFICATION.md	Updated optimizer status, moving Adamax and other implemented optimizers to completed section

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[sydneyrenee]

@copilot resolve the merge markers for this merge of this branch. Rebase and resolve.

[Copilot]

@copilot resolve the merge markers for this merge of this branch. Rebase and resolve.

I've successfully rebased the branch onto origin/main and resolved all merge conflicts.

Changes made:

• Rebased 6 commits onto latest main (b09705c)
• Merged import statements to include both Adamax operations (abs, maximum) and Muon operations (matmul, reshape, transpose)
• Merged optimizer exports to include all optimizers: Adamax, Adagrad, and Muon
• All tests updated to import complete optimizer list

The branch is now up to date with main. The rebased commits are ready on the local branch but need to be force-pushed to update the PR. The merge conflicts have been cleanly resolved with no remaining conflict markers.