diff --git a/index.js b/index.js
deleted file mode 100644
index 1acaa041..00000000
--- a/index.js
+++ /dev/null
@@ -1,52 +0,0 @@
-// The core APIs are in native bindings.
-const core = require(__dirname + '/build/Release/node_mlx.node');
-
-// Helper for creating complex number.
-core.Complex = (re, im) => {
-  return {re, im: im ?? 0};
-};
-
-// The stream helper is in TS.
-core.stream = require('./dist/stream').stream;
-
-// Export core (mx) module.
-exports.core = core;
-
-// Export utils module.
-exports.utils = require('./dist/utils');
-
-// Lazy-load nn/optimizers modules to avoid circular references.
-const cache = {};
-defineLazyModule('nn');
-defineLazyModule('optimizers');
-
-// Export nn/optimizers modules.
-// Note that while it is tempting to get rid of the |cache| object and define
-// the lazy property on |exports| directly, doing so would make the exports
-// undetectable from cjs-module-lexer, and "import {core} from 'mlx'" will not
-// work in Node.js.
-Object.defineProperty(exports, 'nn', {
-  enumerable: true,
-  get() { return cache.nn; }
-});
-Object.defineProperty(exports, 'optimizers', {
-  enumerable: true,
-  get() { return cache.optimizers; }
-});
-
-// Helper to define a lazy loaded property on |cache|.
-function defineLazyModule(name) {
-  Object.defineProperty(cache, name, {
-    configurable: true,
-    enumerable: true,
-    get() {
-      const mod = require(`./dist/${name}`);
-      Object.defineProperty(cache, name, {
-        value: mod,
-        writable: false,
-        configurable: false,
-      });
-      return mod;
-    }
-  });
-}
diff --git a/lib/core.ts b/lib/core.ts
new file mode 100644
index 00000000..bb49f961
--- /dev/null
+++ b/lib/core.ts
@@ -0,0 +1,23 @@
+/// <reference path="../node_mlx.node.d.ts" preserve="true"/>
+import core from '../build/Release/node_mlx.node';
+
+// Helper for creating complex number.
+core.Complex = (re, im) => {
+  return { re, im: im ?? 0 };
+};
+
+// Implementation of the StreamContext.
+core.stream = function stream(s: core.StreamOrDevice): Disposable {
+  const old = core.defaultStream(core.defaultDevice());
+  const target = core.toStream(s);
+  core.setDefaultDevice(target.device);
+  core.setDefaultStream(target);
+  return {
+    [Symbol.dispose]() {
+      core.setDefaultDevice(old.device);
+      core.setDefaultStream(old);
+    },
+  };
+};
+
+export { core, core as mx };
diff --git a/lib/index.ts b/lib/index.ts
new file mode 100644
index 00000000..3b4ed6e3
--- /dev/null
+++ b/lib/index.ts
@@ -0,0 +1,4 @@
+export * from './core';
+export * as nn from './nn';
+export * as optimizers from './optimizers';
+export * as utils from './utils';
diff --git a/lib/nn/index.ts b/lib/nn/index.ts
index cb2d56be..3abe5148 100644
--- a/lib/nn/index.ts
+++ b/lib/nn/index.ts
@@ -1,18 +1,4 @@
-export * from './layers/activations';
-export * from './layers/base';
-export * from './layers/containers';
-export * from './layers/convolution';
-export * from './layers/convolution-transpose';
-export * from './layers/dropout';
-export * from './layers/embedding';
-export * from './layers/linear';
-export * from './layers/normalization';
-export * from './layers/pooling';
-export * from './layers/positional-encoding';
-export * from './layers/quantized';
-export * from './layers/recurrent';
-export * from './layers/transformer';
-export * from './layers/upsample';
+export * from './layers';
 export * from './utils';
 export * as init from './init';
 export * as losses from './losses';
diff --git a/lib/nn/init.ts b/lib/nn/init.ts
index a6203894..12e906bb 100644
--- a/lib/nn/init.ts
+++ b/lib/nn/init.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../..';
+import {core as mx} from '../core';
 
 /**
  * An initializer that returns an array filled with `value`.
diff --git a/lib/nn/layers/activations.ts b/lib/nn/layers/activations.ts
index 31931df4..50e48955 100644
--- a/lib/nn/layers/activations.ts
+++ b/lib/nn/layers/activations.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 
 /**
diff --git a/lib/nn/layers/base.ts b/lib/nn/layers/base.ts
index f26f98b4..6ab76299 100644
--- a/lib/nn/layers/base.ts
+++ b/lib/nn/layers/base.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {
   NestedDict,
   deepEqual,
diff --git a/lib/nn/layers/containers.ts b/lib/nn/layers/containers.ts
index e03bbf98..1fdb97f8 100644
--- a/lib/nn/layers/containers.ts
+++ b/lib/nn/layers/containers.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 
 interface SequentialModule extends Module {
diff --git a/lib/nn/layers/convolution-transpose.ts b/lib/nn/layers/convolution-transpose.ts
index d9e59bf7..31c0b961 100644
--- a/lib/nn/layers/convolution-transpose.ts
+++ b/lib/nn/layers/convolution-transpose.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 
 /**
diff --git a/lib/nn/layers/convolution.ts b/lib/nn/layers/convolution.ts
index 88501b27..b41c2a39 100644
--- a/lib/nn/layers/convolution.ts
+++ b/lib/nn/layers/convolution.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 
 /**
diff --git a/lib/nn/layers/dropout.ts b/lib/nn/layers/dropout.ts
index 1f023d9a..88c120d9 100644
--- a/lib/nn/layers/dropout.ts
+++ b/lib/nn/layers/dropout.ts
@@ -1,4 +1,4 @@
-import {core as mx, utils} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 
 /**
diff --git a/lib/nn/layers/embedding.ts b/lib/nn/layers/embedding.ts
index fd999fc9..be59f20e 100644
--- a/lib/nn/layers/embedding.ts
+++ b/lib/nn/layers/embedding.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 import {QuantizedEmbedding} from './quantized';
 
diff --git a/lib/nn/layers/index.ts b/lib/nn/layers/index.ts
new file mode 100644
index 00000000..4358b2c7
--- /dev/null
+++ b/lib/nn/layers/index.ts
@@ -0,0 +1,15 @@
+export * from './activations';
+export * from './base';
+export * from './containers';
+export * from './convolution';
+export * from './convolution-transpose';
+export * from './dropout';
+export * from './embedding';
+export * from './linear';
+export * from './normalization';
+export * from './pooling';
+export * from './positional-encoding';
+export * from './quantized';
+export * from './recurrent';
+export * from './transformer';
+export * from './upsample';
diff --git a/lib/nn/layers/linear.ts b/lib/nn/layers/linear.ts
index cb5b1c01..e8f32aa6 100644
--- a/lib/nn/layers/linear.ts
+++ b/lib/nn/layers/linear.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 import {QuantizedLinear} from './quantized';
 
diff --git a/lib/nn/layers/normalization.ts b/lib/nn/layers/normalization.ts
index 744daa95..2b1b6d58 100644
--- a/lib/nn/layers/normalization.ts
+++ b/lib/nn/layers/normalization.ts
@@ -1,4 +1,4 @@
-import {core as mx, utils} from '../../..';
+import {core as mx} from '../../core';
 import {range} from './pytools';
 import {Module} from './base';
 
diff --git a/lib/nn/layers/pooling.ts b/lib/nn/layers/pooling.ts
index d37cb9c6..ed077f7a 100644
--- a/lib/nn/layers/pooling.ts
+++ b/lib/nn/layers/pooling.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {accumulate, range} from './pytools';
 import {Module} from './base';
 
diff --git a/lib/nn/layers/positional-encoding.ts b/lib/nn/layers/positional-encoding.ts
index 939b1012..b0c2668a 100644
--- a/lib/nn/layers/positional-encoding.ts
+++ b/lib/nn/layers/positional-encoding.ts
@@ -1,4 +1,4 @@
-import {core as mx, utils} from '../../..';
+import {core as mx} from '../../core';
 import {Module} from './base';
 
 /**
diff --git a/lib/nn/layers/quantized.ts b/lib/nn/layers/quantized.ts
index 8a46cb7a..867a3c63 100644
--- a/lib/nn/layers/quantized.ts
+++ b/lib/nn/layers/quantized.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {toSnakeCase, treeMapWithPath} from '../../utils';
 import {Embedding} from './embedding';
 import {Linear} from './linear';
diff --git a/lib/nn/layers/recurrent.ts b/lib/nn/layers/recurrent.ts
index ce89e9f0..8e4e5ef5 100644
--- a/lib/nn/layers/recurrent.ts
+++ b/lib/nn/layers/recurrent.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {tanh} from './activations';
 import {Module} from './base';
 
diff --git a/lib/nn/layers/transformer.ts b/lib/nn/layers/transformer.ts
index 07ac611d..65fa7735 100644
--- a/lib/nn/layers/transformer.ts
+++ b/lib/nn/layers/transformer.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../../..';
+import {core as mx} from '../../core';
 import {checkpoint} from '../utils';
 import {relu} from './activations';
 import {Module} from './base';
diff --git a/lib/nn/layers/upsample.ts b/lib/nn/layers/upsample.ts
index 793e27d1..eddf32c8 100644
--- a/lib/nn/layers/upsample.ts
+++ b/lib/nn/layers/upsample.ts
@@ -1,4 +1,4 @@
-import {core as mx, utils} from '../../..';
+import {core as mx} from '../../core';
 import {product} from './pytools';
 import {Module} from './base';
 
diff --git a/lib/nn/losses.ts b/lib/nn/losses.ts
index 65fe1101..a50e61da 100644
--- a/lib/nn/losses.ts
+++ b/lib/nn/losses.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../..';
+import {core as mx} from '../core';
 import {deepEqual} from '../utils';
 
 type Reduction = 'none' | 'mean' | 'sum';
@@ -15,7 +15,7 @@ type Reduction = 'none' | 'mean' | 'sum';
  * @param weights - Optional weights for each target. Default: `undefined`.
  * @param axis - The axis over which to compute softmax. Default: `-1`.
  * @param labelSmoothing - Label smoothing factor. Default: `0`.
- * @param reduction - Specifies the reduction to apply to the output: `'none'` | 
+ * @param reduction - Specifies the reduction to apply to the output: `'none'` |
  * `'mean'` | `'sum'`. Default: `'none'`.
  *
  * @returns The computed cross entropy loss.
@@ -370,7 +370,7 @@ export function smoothL1Loss(predictions: mx.array,
  * @param axis - The distribution axis. Default: `-1`.
  * @param p - The norm degree for pairwise distance. Default: `2`.
  * @param margin - Margin for the triplet loss. Defaults to `1.0`.
- * @param eps - Small positive constant to prevent numerical instability. Defaults to `1e-6`. 
+ * @param eps - Small positive constant to prevent numerical instability. Defaults to `1e-6`.
  * @param reduction - Specifies the reduction to apply to the output: `'none'` | `'mean'` | `'sum'`. Default: `'none'`.
  *
  * @returns Computed triplet loss. If reduction is "none", returns a tensor of the same shape as input;
diff --git a/lib/nn/utils.ts b/lib/nn/utils.ts
index 1b716260..b79b91bb 100644
--- a/lib/nn/utils.ts
+++ b/lib/nn/utils.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../..';
+import {core as mx} from '../core';
 import {Module} from './layers/base';
 import {NestedDict} from '../utils';
 
diff --git a/lib/optimizers/optimizers.ts b/lib/optimizers/optimizers.ts
index cbe46f18..ed4869ee 100644
--- a/lib/optimizers/optimizers.ts
+++ b/lib/optimizers/optimizers.ts
@@ -1,4 +1,6 @@
-import {core as mx, nn, utils} from '../..';
+import * as nn from '../nn';
+import * as utils from '../utils';
+import {core as mx} from '../core';
 import {Nested, NestedDict, treeMap} from '../utils';
 
 /**
diff --git a/lib/optimizers/schedulers.ts b/lib/optimizers/schedulers.ts
index 6f46050c..607c9977 100644
--- a/lib/optimizers/schedulers.ts
+++ b/lib/optimizers/schedulers.ts
@@ -1,4 +1,4 @@
-import {core as mx} from '../..';
+import {core as mx} from '../core';
 
 /**
  * Make an exponential decay scheduler.
diff --git a/lib/stream.ts b/lib/stream.ts
deleted file mode 100644
index f0932377..00000000
--- a/lib/stream.ts
+++ /dev/null
@@ -1,16 +0,0 @@
-import mlx from '..';
-
-// Implementation of the StreamContext in JS.
-export function stream(s: mlx.core.StreamOrDevice): Disposable {
-  const mx = require('..').core;  // work around circular import
-  const old = mx.defaultStream(mx.defaultDevice());
-  const target = mx.toStream(s);
-  mx.setDefaultDevice(target.device);
-  mx.setDefaultStream(target);
-  return {
-    [Symbol.dispose]() {
-      mx.setDefaultDevice(old.device);
-      mx.setDefaultStream(old);
-    }
-  };
-}
diff --git a/index.d.ts b/node_mlx.node.d.ts
similarity index 99%
rename from index.d.ts
rename to node_mlx.node.d.ts
index 835ab298..2000c8ae 100644
--- a/index.d.ts
+++ b/node_mlx.node.d.ts
@@ -1,4 +1,4 @@
-export namespace core {
+declare module '*node_mlx.node' {
   // Device.
   type DeviceType = number;
 
@@ -20,7 +20,7 @@ export namespace core {
   }
 
   function defaultStream(device: DeviceOrType): Stream;
-  function setDefaultStream(): void;
+  function setDefaultStream(stream: Stream): void;
   function newStream(device: DeviceOrType): Stream;
   function toStream(s: StreamOrDevice): Stream;
   function stream(s: StreamOrDevice): Disposable;
@@ -475,11 +475,3 @@ export namespace core {
   const newaxis: null;
 }
 
-// The nn module.
-export * as nn from './lib/nn';
-
-// The optim module.
-export * as optimizers from './lib/optimizers';
-
-// The utils module.
-export * as utils from './lib/utils';
diff --git a/package.json b/package.json
index aca5cf9d..611de435 100644
--- a/package.json
+++ b/package.json
@@ -2,8 +2,12 @@
   "name": "@frost-beta/mlx",
   "version": "0.0.1-dev",
   "description": "Node-API bindings for MLX",
-  "main": "index.js",
-  "types": "index.d.ts",
+  "main": "dist/index.js",
+  "types": "dist/index.d.ts",
+  "exports": {
+    ".": "./dist/index.js",
+    "./*": "./dist/*.js"
+  },
   "scripts": {
     "install": "node install.js",
     "prepack": "tsc --build",