Implement mlx.optimizers.Adamax()

docs/OPTIMIZER_API_VERIFICATION.md

@@ -141,12 +141,12 @@ The only remaining work for complete optimizer functionality:
 
 2. **Additional optimizer classes**:
    - ✅ Adam (structure and validation complete, blocked on missing core ops)
+   - ✅ Adamax (fully implemented with all required ops)
+   - ✅ Lion (fully implemented with all required ops)
+   - ✅ RMSprop (fully implemented with all required ops)
    - AdamW
-   - RMSprop
    - Adagrad
    - AdaDelta
-   - Adamax
-   - Lion
    - Adafactor
    - Muon
 

node/CHECKLIST.md

@@ -46,6 +46,7 @@ Reference source map: see `PORTING_PLAN.md` for the list of upstream C++ and Pyt
 - [ ] Complete `SGD.applySingle()` implementation (requires subtract, scalar ops, astype).
 - [x] Add `Adam` optimizer (structure and validation complete, blocked on missing core ops).
 - [ ] Complete `Adam.applySingle()` implementation (requires subtract, divide, square, sqrt, rsqrt, power).
+- [x] Add `Adamax` optimizer (fully implemented with all required ops).
 - [x] Add `Lion` optimizer (fully implemented with all required ops).
 - [x] Add `RMSprop` optimizer (fully implemented with all required ops).
 - [ ] Add `AdamW` optimizer.

node/src/optimizers/index.ts

@@ -6,7 +6,7 @@
  */
 
 import MLXArray, { zeros, zerosLike } from '../core/array';
-import { add, multiply, subtract, sign, square, sqrt, divide, matmul, reshape, transpose } from '../core/ops';
+import { add, multiply, subtract, sign, square, sqrt, divide, abs, maximum, matmul, reshape, transpose } from '../core/ops';
 import { treeMap } from '../utils';
 
 /**
@@ -453,6 +453,92 @@ export class Adam extends Optimizer {
   }
 }
 
+/**
+ * Adamax Optimizer Options
+ */
+export interface AdamaxOptions {
+  /** The learning rate */
+  learningRate: SchedulableParam;
+  /** The coefficients (β₁, β₂) used for computing running averages of gradient and the exponentially weighted infinity norm (default: [0.9, 0.999]) */
+  betas?: [number, number];
+  /** The term ε added to the denominator to improve numerical stability (default: 1e-8) */
+  eps?: number;
+}
+
+/**
+ * The Adamax optimizer, a variant of Adam based on the infinity norm.
+ *
+ * Implements the Adamax algorithm from "Adam: A Method for Stochastic Optimization" (Kingma & Ba, 2015).
+ *
+ * The algorithm updates parameters as follows:
+ *
+ *   m_{t+1} = β₁ * m_t + (1 - β₁) * g_t
+ *   v_{t+1} = max(β₂ * v_t, |g_t|)
+ *   w_{t+1} = w_t - λ * m_{t+1} / (v_{t+1} + ε)
+ *
+ * where λ is the learning rate, m_t is the first moment estimate,
+ * v_t is the exponentially weighted infinity norm, g_t is the gradient,
+ * and ε is a small constant for numerical stability.
+ *
+ * @example
+ * ```typescript
+ * const optimizer = new Adamax({ learningRate: 0.002 });
+ * // ... during training:
+ * // const updatedParams = optimizer.applyGradients(gradients, parameters);
+ * ```
+ */
+export class Adamax extends Adam {
+  constructor(options: AdamaxOptions) {
+    const {
+      learningRate,
+      betas = [0.9, 0.999],
+      eps = 1e-8
+    } = options;
+
+    if (eps < 0.0) {
+      throw new Error(`Epsilon value should be >=0, ${eps} was provided instead`);
+    }
+
+    // Call parent constructor with biasCorrection disabled
+    // Following the original paper, Adamax omits bias correction
+    super({ learningRate, betas, eps, biasCorrection: false });
+  }
+
+  protected initSingle(parameter: MLXArray, state: Record<string, any>): void {
+    // Initialize first moment (m) and exponentially weighted infinity norm (v) with zerosLike
+    state.m = zerosLike(parameter);
+    state.v = zerosLike(parameter);
+  }
+
+  protected applySingle(
+    gradient: MLXArray,
+    parameter: MLXArray,
+    state: Record<string, any>
+  ): MLXArray {
+    const lr = this.learningRate;
+    const [b1, b2] = this.betas;
+    const eps = this.eps;
+
+    // Get moments from state
+    let m = state.m;
+    let v = state.v;
+
+    // Update biased first moment estimate: m = β₁ * m + (1 - β₁) * g
+    m = add(multiply(b1, m), multiply(1 - b1, gradient));
+
+    // Update exponentially weighted infinity norm: v = max(β₂ * v, |g|)
+    v = maximum(multiply(b2, v), abs(gradient));
+
+    // Store updated moments
+    state.m = m;
+    state.v = v;
+
+    // Compute update: w = w - lr * m / (v + ε)
+    const update = divide(multiply(lr, m), add(v, eps));
+    return subtract(parameter, update);
+  }
+}
+
 export interface LionOptions {
   learningRate: SchedulableParam;
   betas?: [number, number];
@@ -865,6 +951,7 @@ export default {
   Optimizer,
   SGD,
   Adam,
+  Adamax,
   Lion,
   Adagrad,
   RMSprop,

node/test/optimizers.test.ts

@@ -1,6 +1,6 @@
 import { describe, it } from 'node:test';
 import assert from 'node:assert';
-import { SGD, Adam, Lion, Adagrad, RMSprop, Muon, Optimizer } from '../src/optimizers';
+import { SGD, Adam, Adamax, Lion, Adagrad, RMSprop, Muon, Optimizer } from '../src/optimizers';
 import { zeros } from '../src/core/array';
 
 describe('mlx.optimizers', () => {
@@ -212,6 +212,89 @@ describe('mlx.optimizers', () => {
     // These should be added once those operations are available
   });
 
+  describe('mlx.optimizers.Adamax', () => {
+    it('should create Adamax optimizer with learning rate', () => {
+      const optimizer = new Adamax({ learningRate: 0.002 });
+      assert.ok(optimizer instanceof Optimizer);
+      assert.ok(optimizer instanceof Adam);
+      assert.ok(optimizer instanceof Adamax);
+      assert.deepStrictEqual(optimizer.betas, [0.9, 0.999]);
+      assert.strictEqual(optimizer.eps, 1e-8);
+      assert.strictEqual(optimizer.biasCorrection, false);
+    });
+
+    it('should create Adamax optimizer with custom betas', () => {
+      const optimizer = new Adamax({
+        learningRate: 0.002,
+        betas: [0.95, 0.9999]
+      });
+      assert.deepStrictEqual(optimizer.betas, [0.95, 0.9999]);
+    });
+
+    it('should create Adamax optimizer with custom epsilon', () => {
+      const optimizer = new Adamax({
+        learningRate: 0.002,
+        eps: 1e-7
+      });
+      assert.strictEqual(optimizer.eps, 1e-7);
+    });
+
+    it('should throw error for negative epsilon', () => {
+      assert.throws(
+        () => new Adamax({ learningRate: 0.002, eps: -1e-8 }),
+        /Epsilon value should be >=0/
+      );
+    });
+
+    it('should have learning rate in state', () => {
+      const optimizer = new Adamax({ learningRate: 0.002 });
+      const lr = optimizer.learningRate;
+      assert.ok(lr);
+      // Check that it's an MLXArray
+      assert.ok(lr.toTypedArray !== undefined);
+    });
+
+    it('should initialize state for parameters', () => {
+      const optimizer = new Adamax({ learningRate: 0.002 });
+      const params = {
+        weight: zeros([3]),
+        bias: zeros([1])
+      };
+
+      optimizer.init(params);
+
+      // Check that state was initialized
+      assert.ok(optimizer.state);
+      assert.ok('step' in optimizer.state);
+      assert.ok('learning_rate' in optimizer.state);
+      assert.ok('weight' in optimizer.state);
+      assert.ok('bias' in optimizer.state);
+
+      // Check that moment estimates were initialized
+      assert.ok('m' in optimizer.state.weight);
+      assert.ok('v' in optimizer.state.weight);
+      assert.ok('m' in optimizer.state.bias);
+      assert.ok('v' in optimizer.state.bias);
+    });
+
+    it('should track step count', () => {
+      const optimizer = new Adamax({ learningRate: 0.002 });
+      const step = optimizer.step;
+      assert.ok(step);
+      assert.strictEqual(step.toTypedArray()[0], 0);
+    });
+
+    it('should allow setting learning rate', () => {
+      const optimizer = new Adamax({ learningRate: 0.002 });
+      optimizer.learningRate = 0.0001;
+      // Note: This test is placeholder since we can't properly set scalar values yet
+      assert.ok(optimizer.learningRate);
+    });
+
+    // Note: Full integration tests with applyGradients would be added here
+    // once we have all operations working properly
+  });
+
   describe('Lion', () => {
     it('should create Lion optimizer with learning rate', () => {
       const optimizer = new Lion({ learningRate: 0.0001 });