Update to MLX 0.23.0

Author: zcbenz

deps/mlx

@@ -45,6 +45,7 @@ export namespace core {
   const int64: Dtype;
   const float16: Dtype;
   const float32: Dtype;
+  const float64: Dtype;
   const bfloat16: Dtype;
   const complex64: Dtype;
 
@@ -201,6 +202,7 @@ export namespace core {
   function atleast3d(...arrays: array[]): array;
   function issubdtype(a: Dtype | DtypeCategory, b: Dtype | DtypeCategory): boolean;
   function bitwiseAnd(a: ScalarOrArray, b: ScalarOrArray, s?: StreamOrDevice): array;
+  function bitwiseInvert(array: ScalarOrArray, s?: StreamOrDevice): array;
   function bitwiseOr(a: ScalarOrArray, b: ScalarOrArray, s?: StreamOrDevice): array;
   function bitwiseXor(a: ScalarOrArray, b: ScalarOrArray, s?: StreamOrDevice): array;
   function broadcastArrays(arrays: array[], s?: StreamOrDevice): array[];
@@ -437,6 +439,10 @@ export namespace core {
     function svd(array: ScalarOrArray, s?: StreamOrDevice): array[];
     function eigvalsh(array: ScalarOrArray, uplo?: string, s?: StreamOrDevice): array[];
     function eigh(array: ScalarOrArray, uplo?: string, s?: StreamOrDevice): array[];
+    function lu(array: ScalarOrArray, s?: StreamOrDevice): array[];
+    function luFactor(array: ScalarOrArray, s?: StreamOrDevice): [array, array];
+    function solve(a: ScalarOrArray, b: ScalarOrArray, s?: StreamOrDevice): array;
+    function solveTriangular(a: ScalarOrArray, b: ScalarOrArray, upper: boolean, s?: StreamOrDevice): array;
   }
 
   // Fast operations.

index.d.ts

@@ -123,7 +123,7 @@ export class Conv2d extends Module {
 
   override toStringExtra(): string {
     return `${this.weight.shape[3]}, ${this.weight.shape[0]}, ` +
-           `kernelSize=${this.weight.shape.slice(1, 2)}, stride=${this.stride}, ` +
+           `kernelSize=${this.weight.shape.slice(1, 3)}, stride=${this.stride}, ` +
            `padding=${this.padding}, dilation=${this.dilation}, ` +
            `groups=${this.groups}, bias=${!!this.bias}`;
   }
@@ -187,7 +187,7 @@ export class Conv3d extends Module {
 
   override toStringExtra(): string {
     return `${this.weight.shape[3]}, ${this.weight.shape[0]}, ` +
-           `kernelSize=${this.weight.shape.slice(1, 3)}, stride=${this.stride}, ` +
+           `kernelSize=${this.weight.shape.slice(1, 4)}, stride=${this.stride}, ` +
            `padding=${this.padding}, dilation=${this.dilation}, ` +
            `bias=${!!this.bias}`;
   }

lib/nn/layers/convolution.ts

@@ -231,6 +231,8 @@ T JsTypedArrayToMxArray(napi_env env, napi_value value) {
       return CreateInstance<T>(static_cast<uint32_t*>(data), shape, mx::uint32);
     case napi_float32_array:
       return CreateInstance<T>(static_cast<float*>(data), shape, mx::float32);
+    case napi_float64_array:
+      return CreateInstance<T>(static_cast<double*>(data), shape, mx::float64);
     default:
       napi_throw_type_error(env, nullptr, "Unsupported TypedArray type.");
       return T();
@@ -330,6 +332,8 @@ auto VisitArrayData(F&& visitor, mx::array* a) {
       return visitor(a->data<mx::float16_t>());
     case mx::float32:
       return visitor(a->data<float>());
+    case mx::float64:
+      return visitor(a->data<double>());
     case mx::bfloat16:
       return visitor(a->data<mx::bfloat16_t>());
     case mx::complex64:
@@ -399,6 +403,7 @@ napi_value ToTypedArray(mx::array* a, napi_env env) {
     case mx::int16: type = napi_int16_array; break;
     case mx::int32: type = napi_int32_array; break;
     case mx::float32: type = napi_float32_array; break;
+    case mx::float64: type = napi_float64_array; break;
     default:
       napi_throw_type_error(env, nullptr, "No matching TypedArray for dtype.");
       return nullptr;
@@ -565,6 +570,8 @@ napi_status Type<mx::Dtype>::ToNode(napi_env env,
     return ConvertToNode(env, &mx::float16, result);
   if (value == mx::float32)
     return ConvertToNode(env, &mx::float32, result);
+  if (value == mx::float64)
+    return ConvertToNode(env, &mx::float64, result);
   if (value == mx::bfloat16)
     return ConvertToNode(env, &mx::bfloat16, result);
   if (value == mx::complex64)
@@ -778,6 +785,7 @@ void InitArray(napi_env env, napi_value exports) {
           "int64", &mx::int64,
           "float16", &mx::float16,
           "float32", &mx::float32,
+          "float64", &mx::float64,
           "bfloat16", &mx::bfloat16,
           "complex64", &mx::complex64);
 

src/array.cc

@@ -47,5 +47,9 @@ void InitLinalg(napi_env env, napi_value exports) {
           "pinv", &mx::linalg::pinv,
           "cross", &linalg::Cross,
           "eigvalsh", &mx::linalg::eigvalsh,
-          "eigh", &mx::linalg::eigvalsh);
+          "eigh", &mx::linalg::eigvalsh,
+          "lu", &mx::linalg::lu,
+          "luFactor", &mx::linalg::lu_factor,
+          "solve", &mx::linalg::solve,
+          "solveTriangular", &mx::linalg::solve_triangular);
 }

src/linalg.cc

@@ -906,6 +906,7 @@ void InitOps(napi_env env, napi_value exports) {
           "atleast3d", NdOpWrapper(&mx::atleast_1d, &mx::atleast_3d),
           "issubdtype", &ops::IsSubDtype,
           "bitwiseAnd", BinOpWrapper(&mx::bitwise_and),
+          "bitwiseInvert", &mx::bitwise_invert,
           "bitwiseOr", BinOpWrapper(&mx::bitwise_or),
           "bitwiseXor", BinOpWrapper(&mx::bitwise_xor),
           "leftShift", BinOpWrapper(&mx::left_shift),

src/ops.cc

@@ -521,4 +521,18 @@ describe('vmap', () => {
     let out = mx.vmap(fun, [0, 1, 1])(a, idx, upd);
     assert.deepEqual(out.shape, [4, 5, 1]);
   });
+
+  it('vmapSplitVmap', () => {
+    const fun = (x: any) => {
+      const [a, b] = mx.split(x, 2, 1);
+      return mx.concatenate([b, a], 1);
+    }
+
+    const x = mx.ones([5, 6, 7]);
+    const y = mx.ones([5, 4, 6, 7]);
+    const fx = fun(x);
+    const fy = mx.vmap(fun, 1)(y);
+    assert.deepEqual(fx.shape, [5, 6, 7]);
+    assert.deepEqual(fy.shape, [4, 5, 6, 7]);
+  });
 });