@genai-fi/nanogpt 0.9.1 → 0.10.0

package/dist/patches/PackedTensor.d.ts

@@ -0,0 +1,12 @@
+import { Rank } from '@tensorflow/tfjs-core/dist/types';
+import { Tensor, Variable } from '@tensorflow/tfjs-core/dist/tensor';
+import { TensorInfo } from '@tensorflow/tfjs-core/dist/tensor_info';
+export interface PackedTensorInfo extends TensorInfo {
+    packed?: boolean;
+}
+export declare class PackableTensor<R extends Rank = Rank> extends Tensor<R> implements PackedTensorInfo {
+    packed: boolean;
+}
+export declare class PackableVariable<R extends Rank = Rank> extends Variable<R> implements PackedTensorInfo {
+    packed: boolean;
+}

package/dist/patches/PackedTensor.js

@@ -0,0 +1,11 @@
+import { T as e, V as s } from "../tensor-DdQUJZlz.js";
+class r extends e {
+  packed = !1;
+}
+class c extends s {
+  packed = !1;
+}
+export {
+  r as PackableTensor,
+  c as PackableVariable
+};

package/dist/patches/engine.d.ts

@@ -0,0 +1,261 @@
+import { BackendTimingInfo, DataMover, KernelBackend } from '@tensorflow/tfjs-core/dist/backends/backend';
+import { Environment } from '@tensorflow/tfjs-core/dist/environment';
+import { NamedAttrMap } from '@tensorflow/tfjs-core/dist/kernel_registry';
+import { TensorInfo, DataId } from '@tensorflow/tfjs-core/dist/tensor_info';
+import { TapeNode } from '@tensorflow/tfjs-core/dist/tape';
+import { DataToGPUOptions, GPUData, Tensor, TensorTracker, Variable } from '@tensorflow/tfjs-core/dist/tensor';
+import { GradSaveFunc, NamedTensorMap, NamedVariableMap, TensorContainer } from '@tensorflow/tfjs-core/dist/tensor_types';
+import { BackendValues, DataType, DataValues } from '@tensorflow/tfjs-core/dist/types';
+/**
+ * A function that computes an output. The save function is for saving tensors
+ * computed in the forward pass, that we need in the backward pass.
+ */
+export type ForwardFunc<T> = (backend: KernelBackend, save?: GradSaveFunc) => T;
+/**
+ * @docalias (a: Tensor, b: Tensor,..., save?: Function) => {
+ *   value: Tensor,
+ *   gradFunc: (dy: Tensor, saved?: NamedTensorMap) => Tensor | Tensor[]
+ * }
+ */
+export type CustomGradientFunc<T extends Tensor> = (...inputs: Array<Tensor | GradSaveFunc>) => {
+    value: T;
+    gradFunc: (dy: T, saved: Tensor[]) => Tensor | Tensor[];
+};
+export type MemoryInfo = {
+    numTensors: number;
+    numDataBuffers: number;
+    numBytes: number;
+    unreliable?: boolean;
+    reasons: string[];
+};
+type KernelInfo = {
+    name: string;
+    bytesAdded: number;
+    totalBytesSnapshot: number;
+    tensorsAdded: number;
+    totalTensorsSnapshot: number;
+    inputShapes: number[][];
+    outputShapes: number[][];
+    kernelTimeMs: number | {
+        error: string;
+    } | Promise<number | {
+        error: string;
+    }>;
+    extraInfo: string | Promise<string>;
+};
+export type ProfileInfo = {
+    newBytes: number;
+    newTensors: number;
+    peakBytes: number;
+    kernels: KernelInfo[];
+    result: TensorContainer;
+    kernelNames: string[];
+};
+export interface TimingInfo extends BackendTimingInfo {
+    wallMs: number;
+}
+/** @docalias Function */
+export type ScopeFn<T extends TensorContainer> = () => T;
+interface ScopeState {
+    track: Tensor[];
+    name: string;
+    id: number;
+}
+declare class EngineState {
+    registeredVariables: NamedVariableMap;
+    nextTapeNodeId: number;
+    numBytes: number;
+    numTensors: number;
+    numStringTensors: number;
+    numDataBuffers: number;
+    activeTape: TapeNode[];
+    gradientDepth: number;
+    kernelDepth: number;
+    activeScope: ScopeState;
+    scopeStack: ScopeState[];
+    /**
+     * Keeps track of the number of data moves during a kernel execution. We
+     * maintain a stack since kernels can call other kernels, recursively.
+     */
+    numDataMovesStack: number[];
+    nextScopeId: number;
+    tensorInfo: WeakMap<object, {
+        backend: KernelBackend;
+        bytes: number;
+        dtype: DataType;
+        shape: number[];
+    }>;
+    profiling: boolean;
+    activeProfile: ProfileInfo;
+    dispose(): void;
+}
+export declare class Engine implements TensorTracker, DataMover {
+    ENV: Environment;
+    version: string;
+    state: EngineState;
+    backendName: string;
+    registry: {
+        [id: string]: KernelBackend;
+    };
+    registryFactory: {
+        [id: string]: {
+            factory: () => KernelBackend | Promise<KernelBackend>;
+            priority: number;
+        };
+    };
+    private profiler;
+    private backendInstance;
+    private pendingBackendInit;
+    private pendingBackendInitId;
+    constructor(ENV: Environment);
+    ready(): Promise<void>;
+    get backend(): KernelBackend;
+    backendNames(): string[];
+    findBackend(backendName: string): KernelBackend;
+    findBackendFactory(backendName: string): () => KernelBackend | Promise<KernelBackend>;
+    registerBackend(backendName: string, factory: () => KernelBackend | Promise<KernelBackend>, priority?: number): boolean;
+    setBackend(backendName: string): Promise<boolean>;
+    private setupRegisteredKernels;
+    private disposeRegisteredKernels;
+    /**
+     * Initializes a backend by looking up the backend name in the factory
+     * registry and calling the factory method. Returns a boolean representing
+     * whether the initialization of the backend succeeded. Throws an error if
+     * there is no backend in the factory registry.
+     */
+    private initializeBackend;
+    removeBackend(backendName: string): void;
+    private getSortedBackends;
+    private initializeBackendsAndReturnBest;
+    moveData(backend: KernelBackend, dataId: DataId): void;
+    tidy<T extends TensorContainer>(nameOrFn: string | ScopeFn<T>, fn?: ScopeFn<T>): T;
+    private scopedRun;
+    private static nextTensorId;
+    private nextTensorId;
+    private static nextVariableId;
+    private nextVariableId;
+    /**
+     * This method is called instead of the public-facing tensor.clone() when
+     * saving a tensor for backwards pass. It makes sure to add the clone
+     * operation to the tape regardless of being called inside a kernel
+     * execution.
+     */
+    private clone;
+    /**
+     * Execute a kernel with the given name and return the output tensor.
+     *
+     * @param kernelName The name of the kernel to execute.
+     * @param inputs A map of input names to tensors.
+     * @param attrs A map of attribute names to their values. An attribute is a
+     *     primitive (non-tensor) input to the kernel.
+     * @param inputsToSave A list of tensors, inputs to save for the backprop
+     *     computation.
+     * @param outputsToSave A list of booleans, specifying which output to save
+     *     for the backprop computation. These are booleans since the output
+     * tensors are not visible to the user.
+     */
+    runKernel<T extends Tensor | Tensor[]>(kernelName: string, inputs: NamedTensorMap, attrs?: NamedAttrMap): T;
+    private shouldCheckForMemLeaks;
+    private checkKernelForMemLeak;
+    /**
+     * Internal helper method to execute a kernel Func
+     *
+     * Use `runKernel` to execute kernels from outside of engine.
+     */
+    private runKernelFunc;
+    /**
+     * Saves tensors used in forward mode for use in backward mode.
+     *
+     * @param tensors the list of tensors to save.
+     */
+    private saveTensorsForBackwardMode;
+    /**
+     * Returns a list of tensors to save for a given gradient calculation.
+     *
+     * @param kernelName name of kernel to look up gradient for.
+     * @param inputs a map of input tensors.
+     * @param outputs an array of output tensors from forward mode of kernel.
+     */
+    private getTensorsForGradient;
+    /**
+     * Internal method used by public APIs for tensor creation. Makes a new
+     * tensor with the provided shape, dtype and values. It always
+     * creates a new data id and writes the values to the underlying backend.
+     */
+    makeTensor(values: DataValues, shape: number[], dtype: DataType, backend?: KernelBackend): Tensor;
+    /**
+     * Internal method used by backends. Makes a new tensor
+     * that is a wrapper around an existing data id. It doesn't create
+     * a new data id, only increments the ref count used in memory tracking.
+     * @deprecated
+     */
+    makeTensorFromDataId(dataId: DataId, shape: number[], dtype: DataType, backend?: KernelBackend): Tensor;
+    /**
+     * Internal method used by backends. Makes a new tensor that is a wrapper
+     * around an existing data id in TensorInfo. It doesn't create a new data id,
+     * only increments the ref count used in memory tracking.
+     */
+    makeTensorFromTensorInfo(tensorInfo: TensorInfo, backend?: KernelBackend): Tensor;
+    makeVariable(initialValue: Tensor, trainable?: boolean, name?: string, dtype?: DataType): Variable;
+    trackTensor(a: Tensor, backend: KernelBackend): void;
+    incRef(a: Tensor, backend: KernelBackend): void;
+    removeDataId(dataId: DataId, backend: KernelBackend): void;
+    disposeTensor(a: Tensor): void;
+    disposeVariables(): void;
+    disposeVariable(v: Variable): void;
+    memory(): MemoryInfo;
+    profile(query: () => TensorContainer | Promise<TensorContainer>): Promise<ProfileInfo>;
+    isTapeOn(): boolean;
+    private addTapeNode;
+    keep<T extends Tensor>(result: T): T;
+    private startTape;
+    private endTape;
+    /**
+     * Start a scope. Use this with endScope() to achieve the same functionality
+     * as scope() without the need for a function closure.
+     */
+    startScope(name?: string): void;
+    /**
+     * End a scope. Use this with startScope() to achieve the same functionality
+     * as scope() without the need for a function closure.
+     */
+    endScope(result?: TensorContainer): void;
+    /**
+     * Returns gradients of `f` with respect to each of the `xs`. The gradients
+     * returned are of the same length as `xs`, but some might be null if `f`
+     * was not a function of that `x`. It also takes optional dy to multiply the
+     * gradient, which defaults to `1`.
+     */
+    gradients<T extends Tensor>(f: () => T, xs: Tensor[], dy?: T, allowNoGradients?: boolean): {
+        value: T;
+        grads: Tensor[];
+    };
+    customGrad<T extends Tensor>(f: CustomGradientFunc<T>): (...args: Array<Tensor>) => T;
+    readSync(dataId: DataId): BackendValues;
+    read(dataId: DataId): Promise<BackendValues>;
+    readToGPU(dataId: DataId, options?: DataToGPUOptions): GPUData;
+    time(query: () => void): Promise<TimingInfo>;
+    /**
+     * Tracks a Tensor in the current scope to be automatically cleaned up
+     * when the current scope ends, and returns the value.
+     *
+     * @param result The Tensor to track in the current scope.
+     */
+    private track;
+    get registeredVariables(): NamedVariableMap;
+    /**
+     * Resets the engine state. Removes all backends but does not remove
+     * registered backend factories.
+     */
+    reset(): void;
+}
+export declare function getOrMakeEngine(): Engine;
+export declare const ENGINE: Engine;
+/**
+ * A implementation of the add op for use within engine and tape.
+ *
+ * This allows us to avoid a circular dependency between add.ts and engine.
+ * It is exported to be available in tape tests.
+ */
+export declare function add(a: Tensor, b: Tensor): Tensor;
+export {};

package/dist/patches/engine.js

@@ -0,0 +1,10 @@
+import "../tensor_util-DV-FP5Q3.js";
+import "../tensor-DdQUJZlz.js";
+import "./PackedTensor.js";
+import { a as t, E, c as g, g as m } from "../utilities/packed.js";
+export {
+  t as ENGINE,
+  E as Engine,
+  g as add,
+  m as getOrMakeEngine
+};

package/dist/patches/tape.d.ts

@@ -0,0 +1,12 @@
+import { Tensor } from '@tensorflow/tfjs-core/dist/tensor';
+import { TapeNode } from '@tensorflow/tfjs-core/dist/tape';
+/**
+ * Backpropagate gradients through the filtered TapeNodes.
+ *
+ * @param tensorAccumulatedGradientMap A map of Tensor to its gradient. This map
+ * is mutated by this method.
+ * @param filteredTape The filtered TapeNodes to backprop through.
+ */
+export declare function backpropagateGradients(tensorAccumulatedGradientMap: {
+    [tensorId: number]: Tensor;
+}, filteredTape: TapeNode[], tidy: (f: Function) => Tensor, add: (a: Tensor, b: Tensor) => Tensor): void;

package/dist/patches/tape.js

@@ -0,0 +1,5 @@
+import { b as p } from "../utilities/packed.js";
+import "../tensor-DdQUJZlz.js";
+export {
+  p as backpropagateGradients
+};

package/dist/patches/webgpu_backend.d.ts

@@ -0,0 +1,18 @@
+import { DataType, TensorInfo } from '@tensorflow/tfjs-core';
+import { WebGPUBackend } from '@tensorflow/tfjs-backend-webgpu/dist/webgpu';
+import { WebGPUProgram } from './webgpu_program';
+type ProgramUniform = Array<{
+    type: string;
+    data: number[];
+}>;
+interface ExtendedAdapterInfo extends GPUAdapterInfo {
+    subgroupMaxSize?: number;
+    subgroupMinSize?: number;
+}
+export default class WebGPUBackendPatch extends WebGPUBackend {
+    readonly subgroupMaxSize: number;
+    readonly subgroupMinSize: number;
+    constructor(device: GPUDevice, adapterInfo?: ExtendedAdapterInfo);
+    runWebGPUProgram(program: WebGPUProgram, inputs: TensorInfo[], outputDtype: DataType, programDefinedUniform?: ProgramUniform, output?: TensorInfo): TensorInfo;
+}
+export {};

package/dist/patches/webgpu_backend.js

@@ -0,0 +1,57 @@
+import "../index-ZyQhjEPo.js";
+import "../webgpu_util-BBCnKm2X.js";
+import { W as c } from "../backend_webgpu-B0u2ndUn.js";
+import { compileProgram as l } from "./webgpu_program.js";
+import { m } from "../webgpu_program-Cigz-7RF.js";
+import { s as P, g as y, e as M, a as n } from "../tensor-DdQUJZlz.js";
+const b = (h, s) => {
+  const i = h.limits.maxComputeWorkgroupsPerDimension, t = s.dispatchLayout, a = s.dispatch;
+  if (a.every((r) => r <= i))
+    return a;
+  n(
+    a[0] > i && t.y === void 0 && t.z === void 0,
+    () => "Dispatch size exceeds WebGPU limits in Y or Z dimension."
+  );
+  let e = Math.ceil(Math.sqrt(a[0]));
+  return e > i ? (e = Math.ceil(Math.cbrt(a[0])), n(
+    e <= i,
+    () => "Total dispatch size exceeds WebGPU maximum."
+  ), [e, e, e]) : [e, e, 1];
+};
+class U extends c {
+  subgroupMaxSize;
+  subgroupMinSize;
+  constructor(s, i) {
+    super(s, i), this.subgroupMaxSize = i?.subgroupMaxSize ?? 0, this.subgroupMinSize = i?.subgroupMinSize ?? 0;
+  }
+  runWebGPUProgram(s, i, t, a, e) {
+    if (e || (e = this.makeTensorInfo(s.outputShape, t)), P(e.shape) === 0)
+      return this.tensorMap.get(e.dataId).values = y(e.dtype, 0), e;
+    this.uploadToGPU(e.dataId), s.dispatch = b(this.device, s);
+    const r = i.map((o, p) => {
+      if (o.dtype === "complex64")
+        throw new Error(
+          "GPGPUProgram does not support complex64 input. For complex64 dtypes, please separate the program into real and imaginary parts."
+        );
+      return this.uploadToGPU(o.dataId), {
+        // Returning dtype from tensorMap because it reflects dtype
+        // of underlying buffer, rather than abstract dtype.
+        dtype: this.tensorMap.get(o.dataId).dtype,
+        shape: o.shape,
+        name: s.variableNames[p]
+      };
+    });
+    s.shaderKey = m(s, r, e);
+    const d = M().getBool("WEBGPU_ENGINE_COMPILE_ONLY");
+    return s.shaderKey in this.pipelineCache || (this.pipelineCache[s.shaderKey] = l(
+      this.device,
+      s,
+      r,
+      e,
+      d
+    )), s.pipeline = this.pipelineCache[s.shaderKey], d || this.recordAndSubmit(s, e, i, a), e;
+  }
+}
+export {
+  U as default
+};

package/dist/{tensor-CZr4dh61.js → patches/webgpu_base.d.ts}

@@ -1,7 +1,6 @@
-import { C as t, D as a } from "./index-BzFyqcy-.js";
 /**
  * @license
- * Copyright 2018 Google LLC. All Rights Reserved.
+ * Copyright 2022 Google Inc. All Rights Reserved.
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
  * You may obtain a copy of the License at
@@ -15,10 +14,8 @@ import { C as t, D as a } from "./index-BzFyqcy-.js";
  * limitations under the License.
  * =============================================================================
  */
-function f(r, n, e) {
-  const o = t(r, e);
-  return a(r, n, o, e);
+export interface GPUOptions {
+    powerPreference?: 'low-power' | 'high-performance';
+    disableSubgroups?: boolean;
 }
-export {
-  f as t
-};
+export declare function registerWebGPUBackend(options?: GPUOptions): void;

package/dist/patches/webgpu_base.js

@@ -0,0 +1,34 @@
+import { f as n } from "../index-ZyQhjEPo.js";
+import p from "./webgpu_backend.js";
+function c(a) {
+  n(
+    "webgpu",
+    async () => {
+      const i = {
+        powerPreference: a?.powerPreference ?? "high-performance"
+      };
+      console.log("Using custom WebGPU backend with power preference:", i.powerPreference);
+      const e = await navigator.gpu.requestAdapter(i), t = {}, o = [];
+      e.features.has("timestamp-query") && o.push("timestamp-query"), e.features.has("bgra8unorm-storage") && o.push(["bgra8unorm-storage"]), !a?.disableSubgroups && e.features.has("subgroups") && o.push("subgroups"), t.requiredFeatures = o;
+      const r = e.limits;
+      t.requiredLimits = {
+        maxComputeWorkgroupStorageSize: r.maxComputeWorkgroupStorageSize,
+        maxComputeWorkgroupsPerDimension: r.maxComputeWorkgroupsPerDimension,
+        maxStorageBufferBindingSize: r.maxStorageBufferBindingSize,
+        maxBufferSize: r.maxBufferSize,
+        maxComputeWorkgroupSizeX: r.maxComputeWorkgroupSizeX,
+        maxComputeInvocationsPerWorkgroup: r.maxComputeInvocationsPerWorkgroup
+      };
+      const u = await e.requestDevice(t), s = "info" in e ? e.info : "requestAdapterInfo" in e ? (
+        // eslint-disable-next-line @typescript-eslint/no-explicit-any
+        await e.requestAdapterInfo()
+      ) : void 0;
+      return new p(u, s);
+    },
+    3
+    /*priority*/
+  );
+}
+export {
+  c as registerWebGPUBackend
+};

package/dist/patches/webgpu_program.d.ts

@@ -0,0 +1,36 @@
+import { DataType, TensorInfo } from '@tensorflow/tfjs-core';
+export declare enum PixelsOpType {
+    FROM_PIXELS = 0,
+    DRAW = 1
+}
+export interface WebGPUProgram {
+    atomic?: boolean;
+    subgroups?: boolean;
+    subgroupBuiltins?: boolean;
+    dispatch: [number, number, number];
+    dispatchLayout: {
+        x: number[];
+        y?: number[];
+        z?: number[];
+    };
+    outputComponent?: number;
+    outputShape: number[];
+    pixelsOpType?: PixelsOpType;
+    shaderKey: string;
+    size?: boolean;
+    uniforms?: string;
+    variableNames: string[];
+    variableComponents?: number[];
+    workgroupSize: [number, number, number];
+    workPerThread?: number;
+    pipeline?: GPUComputePipeline | Promise<GPUComputePipeline>;
+    getUserCode: () => string;
+}
+export declare const compileProgram: (device: GPUDevice, program: WebGPUProgram, inputsData: InputInfo[], output: TensorInfo, parallelCompilation: boolean) => GPUComputePipeline | Promise<GPUComputePipeline>;
+type InputInfo = {
+    dtype: DataType;
+    shape: number[];
+    name: string;
+};
+export declare function getStartHeaderString(useGlobalIndex: boolean, program: WebGPUProgram): string;
+export {};