Implement TypeScript wrapper for mlx.core.arange()

[Copilot]

Overview

This PR implements the TypeScript/JavaScript wrapper for mlx.core.arange(), making the function accessible from the Node.js API. The C++ implementation was already present in node/src/native/array.cc and registered in the addon, but lacked the TypeScript interface layer.

Changes

TypeScript Wrapper (node/src/core/ops.ts)

Added the arange() function with full signature support:

// Single argument: arange(stop)
const arr1 = arange(10);  // [0, 1, 2, ..., 9]

// Two arguments: arange(start, stop)
const arr2 = arange(5, 10);  // [5, 6, 7, 8, 9]

// Three arguments: arange(start, stop, step)
const arr3 = arange(0, 10, 2);  // [0, 2, 4, 6, 8]

// With optional dtype and stream
const arr4 = arange(10, undefined, undefined, { dtype: float32 });

The implementation:

• Handles all three calling patterns with proper argument parsing
• Supports optional dtype parameter for explicit type control
• Supports optional stream parameter for stream-aware execution
• Includes comprehensive JSDoc documentation with examples
• Follows existing patterns from zeros, ones, and full operations

Exports

• node/src/core/index.ts: Exported arange and ArangeOptions type from core module
• node/src/index.ts: Added top-level exports for arange and dtype constants (int32, float32, float16, etc.) for convenience

Comprehensive Tests (node/test/core/ops.test.ts)

Added 17 test cases covering:

• Basic range generation (1, 2, and 3 arguments)
• Fractional steps: arange(0, 3, 0.5) → [0, 0.5, 1, 1.5, 2, 2.5]
• Negative steps: arange(0, -5, -1) → [0, -1, -2, -3, -4]
• Edge cases: empty ranges, step larger than range
• Dtype inference: integers default to int32, floats to float32
• Explicit dtype parameters: float16, float32, int64, uint32
• Stream support for async execution

Tests align with the C++ test suite in tests/creations_tests.cpp.

Implementation Notes

• Dtype Inference: The C++ backend automatically handles dtype inference matching Python MLX behavior
• Type Safety: Added ArangeOptions interface for TypeScript users
• Minimal Changes: Only added the necessary TypeScript wrapper layer—no C++ modifications needed
• Consistency: Follows the same patterns as other factory functions in the codebase
• Zero Breaking Changes: Pure additive change, no existing functionality modified

Testing

TypeScript compilation passes successfully. The implementation is ready for integration testing once the native addon build is available in CI.

Fixes #[issue_number]

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Implement mlx.core.arange()</issue_title>
<issue_description>## 🎯 Implement mlx.core.arange()

Priority: medium | Module: core | Type: other | Category: core

---

📋 Quick Reference

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

---

🚀 Step-by-Step Implementation

Step 1: Review Python Implementation

# See the Python binding
grep -B 5 -A 30 '"arange"' python/src/ops.cpp

Step 2: Implement in Node.js

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

Napi::Value Arange(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/src/ops.cpp 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::arange(/* 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("arange failed: ") + e.what())
        .ThrowAsJavaScriptException();
    return env.Null();
  }
}

Step 3: Register the Function

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

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

Step 4: Add Tests

File: node/test/ops.test.js

const mx = require('..');

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

---

📚 Resources

Similar Implementations

Look at these in node/src/native/array.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/array.cc
• Registered in Init()
• Added tests in node/test/ops.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 #504

---

✨ 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 TypeScript wrapper for mlx.core.arange(), making the function accessible from the Node.js API. The implementation adds the JavaScript/TypeScript interface layer over the existing C++ implementation.

• Added TypeScript wrapper function with comprehensive JSDoc documentation and overloaded signatures
• Exported the function and associated types from the core module
• Added extensive test coverage for all usage patterns and edge cases

Reviewed Changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 2 comments.

File	Description
node/src/core/ops.ts	Implements the arange function with TypeScript wrapper and ArangeOptions interface
node/src/core/index.ts	Exports arange function and ArangeOptions type from core module
node/src/index.ts	Adds top-level exports for arange and dtype constants
node/test/core/ops.test.ts	Comprehensive test suite covering all arange functionality

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[sydneyrenee]

@copilot Address comments

[Copilot]

@copilot Address comments

I don't see any specific comments or feedback in the PR that need to be addressed. The implementation appears complete:

• ✅ TypeScript wrapper for arange() is implemented with proper overloading
• ✅ 17 comprehensive test cases added
• ✅ Exports configured in both core/index.ts and top-level index.ts
• ✅ TypeScript compilation passes with no errors
• ✅ JSDoc documentation has been updated (commit 91db9fd)

Could you please point me to the specific comments or issues you'd like me to address?