Set noImplicitAny in tsconfig

lib/index.d.ts

@@ -349,7 +349,7 @@ export namespace core {
 
   // Memory management.
   function tidy<U>(func: () => U): U;
-  function dispose(tree: unknown);
+  function dispose(tree: unknown): void;
   function getWrappersCount(): number;
 
   // Metal.

lib/nn/layers/base.ts

@@ -341,7 +341,7 @@ export abstract class Module {
    * @returns The module instance after updating the parameters.
    */
   update(parameters: NestedDict<mx.array>): this {
-    const apply = (target: object, parameters: NestedDict<mx.array>) => {
+    const apply = (target: Record<string, any>, parameters: NestedDict<mx.array>) => {
       if (typeof parameters !== 'object')
         return;
       for (let k in parameters) {
@@ -400,7 +400,7 @@ export abstract class Module {
    * @returns The module instance after updating the submodules.
    */
   updateModules(modules: NestedDict<Module>): this {
-    const apply = (target: object, modules: object) => {
+    const apply = (target: Record<string, any>, modules: Record<string, any>) => {
       if (typeof modules !== 'object')
         return;
       for (let k in modules) {
@@ -602,7 +602,7 @@ export abstract class Module {
 }
 
 // Helpers.
-function defaultIsLeafFn(m, k, v) {
+function defaultIsLeafFn(m: Module, k: string, v: unknown) {
   if (typeof v !== 'object')
     return true;
   return !Array.isArray(v) && !isDict(v);
@@ -643,9 +643,10 @@ function unwrap(model: Module,
   if (typeof value === 'object') {
     const newValue: Record<string, unknown> = {};
     for (let k in value) {
-      let key = `${valueKey}.${k}`;
-      if (filterFn(model, key, value[k])) {
-        newValue[k] = unwrap(model, key, value[k], filterFn, mapFn, isLeafFn);
+      const key = `${valueKey}.${k}`;
+      const v = value[k as keyof typeof value];
+      if (filterFn(model, key, v)) {
+        newValue[k] = unwrap(model, key, v, filterFn, mapFn, isLeafFn);
       } else {
         newValue[k] = {};
       }

lib/nn/layers/normalization.ts

@@ -322,7 +322,7 @@ export class BatchNorm extends Module {
     }
   }
 
-  override unfreeze(...args): this {
+  override unfreeze(...args: any[]): this {
     super.unfreeze(...args);
     this.freeze(false, ['runningMean', 'runningVar']);
     return this;

lib/nn/layers/positional-encoding.ts

@@ -137,20 +137,20 @@ export class ALiBi extends Module {
     return ALiBi.mask;
   }
 
-  static createAlibiSlope(numHeads): mx.array {
+  static createAlibiSlope(numHeads: number): mx.array {
     const x = Math.pow(Math.pow(2, 8), 1 / numHeads);
     const out = mx.power(x, mx.negative(mx.arange(1, numHeads + 1)));
     return mx.expandDims(out, [-1, -2]);
   }
 
-  forward(attentionScores: mx.array, offset = 0, mask = null): mx.array {
+  forward(attentionScores: mx.array, offset = 0, mask?: mx.array): mx.array {
     let alibiMask = ALiBi.createAlibiMatrix(
       attentionScores.shape[attentionScores.shape.length - 2] + offset,
       attentionScores.shape[attentionScores.shape.length - 1],
       attentionScores.shape[1],
       offset,
       attentionScores.dtype);
-    if (mask !== null) {
+    if (mask) {
       alibiMask = mx.add(alibiMask, mask);
     }
     return mx.add(attentionScores, alibiMask);

lib/nn/layers/quantized.ts

@@ -211,7 +211,7 @@ export class QuantizedLinear extends Module {
 
   // Wrap unfreeze so that we unfreeze any layers we might contain but
   // our parameters will remain frozen.
-  override unfreeze(...args): this {
+  override unfreeze(...args: any[]): this {
     super.unfreeze(...args);
     return this.freeze(false);
   }

lib/nn/layers/upsample.ts

@@ -216,7 +216,7 @@ export function upsampleNearest(x: mx.array,
 
 function interpolate(x: mx.array,
                      scaleFactor: number[],
-                     indicesFn,
+                     indicesFn: (N: number, scale: number, alignCorners: boolean, dim: number, ndims: number) => [mx.array, mx.array][],
                      alignCorners = false): mx.array {
   const dims = x.ndim - 2;
   if (dims !== scaleFactor.length) {

lib/optimizers/optimizers.ts

@@ -54,7 +54,7 @@ export abstract class Optimizer {
    */
   init(parameters: Nested<mx.array>): void {
     // Iniatilize the optimizer state to match the parameter state.
-    function updateState(params: Nested<mx.array>, state: Nested<unknown>) {
+    function updateState(params: Nested<mx.array>, state: Record<string, any>) {
       if (Array.isArray(params) && Array.isArray(state)) {
         for (let i = 0; i < state.length; ++i)
           state[i] = updateState(params[i], state[i]);

lib/utils.ts

@@ -32,8 +32,8 @@
  * @returns An Object with the new values returned by `fn`.
  */
 export function treeMap(func: (...args: any[]) => unknown,
-                        tree: unknown,
-                        rest?: object[],
+                        tree: Record<string, any>,
+                        rest?: Record<string, any>[],
                         isLeaf?: (node: unknown) => boolean): unknown {
   if (isLeaf && isLeaf(tree)) {
     if (rest)
@@ -43,7 +43,7 @@ export function treeMap(func: (...args: any[]) => unknown,
   } else if (Array.isArray(tree)) {
     return tree.map((child, i) => treeMap(func, child, rest?.map(r => r[i]), isLeaf));
   } else if (typeof tree === 'object' && isDict(tree)) {
-    const newTree = {};
+    const newTree: Record<string, any> = {};
     for (const k in tree) {
       newTree[k] = treeMap(func, tree[k], rest?.map(r => r[k]), isLeaf);
     }
@@ -84,8 +84,8 @@ export function treeMap(func: (...args: any[]) => unknown,
  * ```
  */
 export function treeMapWithPath(fn: (path?: string, ...args: unknown[]) => unknown,
-                                tree: unknown,
-                                rest?: object[],
+                                tree: Record<string, any>,
+                                rest?: Record<string, any>[],
                                 isLeaf?: (node: unknown) => boolean,
                                 path?: string): unknown {
   if (isLeaf && isLeaf(tree)) {
@@ -96,7 +96,7 @@ export function treeMapWithPath(fn: (path?: string, ...args: unknown[]) => unkno
   } else if (Array.isArray(tree)) {
     return tree.map((child, i) => treeMapWithPath(fn, child, rest?.map(r => r[i]), isLeaf, path ? `${path}.${i}` : `${i}`));
   } else if (typeof tree === 'object' && isDict(tree)) {
-    const newTree = {};
+    const newTree: Record<string, any> = {};
     for (const k in tree) {
       newTree[k] = treeMapWithPath(fn, tree[k], rest?.map(r => r[k]), isLeaf, path ? `${path}.${k}` : `${k}`);
     }
@@ -138,7 +138,7 @@ export function treeMapWithPath(fn: (path?: string, ...args: unknown[]) => unkno
  * // Outputs: [['hello.0.0.0', 0]]
  * ```
  */
-export function treeFlatten(tree: unknown,
+export function treeFlatten(tree: Record<string, any>,
                             prefix: string = '',
                             isLeaf?: (node: unknown) => boolean,
                             convertKey?: (key: string) => string): [string, unknown][] {
@@ -199,12 +199,12 @@ export function treeUnflatten(tree: [string, unknown][],
   // Recursively map them to the original container.
   if (isList) {
     const keys = Object.keys(children).sort().map((idx) => [ parseInt(idx), idx ]);
-    const newList = [];
+    const newList: any[] = [];
     for (const [i, k] of keys)
-      newList[i] = treeUnflatten(children[k], convertKey);
+      newList[i as number] = treeUnflatten(children[k], convertKey);
     return newList;
   } else {
-    const newTree = {};
+    const newTree: Record<string, any> = {};
     for (let k in children) {
       const key = convertKey ? convertKey(k) : k;
       newTree[key] = treeUnflatten(children[k], convertKey);

tsconfig.json

@@ -3,6 +3,7 @@
     "outDir": "dist",
     "rootDir": "lib",
     "noImplicitOverride": true,
+    "noImplicitAny": true,
     "module": "commonjs",
     "esModuleInterop": true,
     "target": "es2022",