import '@tensorflow/tfjs-backend-webgl'; import '@tensorflow/tfjs-backend-webgpu'; import '@tensorflow/tfjs-backend-cpu'; import * as wasm from '@tensorflow/tfjs-backend-wasm'; import * as tf from '@tensorflow/tfjs-core'; import { MLComputeResult, MLGraph, MLNamedArrayBufferViews } from './graph'; import * as utils from './utils'; /** @internal */ export enum MLContextType { 'default' = 'default', 'webgpu' = 'webgpu' } /** * [spec](https://webmachinelearning.github.io/webnn/#enumdef-mldevicetype) */ export enum MLDeviceType { 'cpu' = 'cpu', 'gpu' = 'gpu' } /** * [spec](https://webmachinelearning.github.io/webnn/#enumdef-mlpowerpreference) */ export enum MLPowerPreference { 'default' = 'default', 'high-performance' = 'high-performance', 'low-power' = 'low-power' } /** * [spec](https://webmachinelearning.github.io/webnn/#dictdef-mlcontextoptions) */ export interface MLContextOptions { deviceType?: MLDeviceType; powerPreference?: MLPowerPreference; } /** * [API spec](https://webmachinelearning.github.io/webnn/#api-mlcontext) */ export class MLContext { private options_: MLContextOptions; private type_: MLContextType; /** @internal */ // eslint-disable-next-line @typescript-eslint/no-empty-function constructor(options: MLContextOptions = {}) { utils.assert(options instanceof Object, 'Invalid options.'); if (options.deviceType !== undefined) { utils.assert( options.deviceType in MLDeviceType, 'Invalid device type.'); } if (options.powerPreference !== undefined) { utils.assert( options.powerPreference in MLPowerPreference, 'Invalid power preference.'); } this.options_ = options; this.type_ = MLContextType.default; } /** @internal */ get options(): MLContextOptions { return this.options_; } /** @internal */ get type(): MLContextType { return this.type_; } /** * [spec](https://webmachinelearning.github.io/webnn/#api-mlcontext-compute) */ async compute( graph: MLGraph, inputs: MLNamedArrayBufferViews, outputs: MLNamedArrayBufferViews): Promise { const result = await graph.compute(inputs, outputs); return result; } /** @internal */ // Expose tf interfance for setting backend. get tf(): unknown { return tf; } /** @internal */ // Expose wasm interface for supporting configure threads for wasm backend. // wasm.setThreadsCount(n) get wasm(): unknown { return wasm; } }