torch.register_forward
function register_forward<Op extends OpName, DT extends DType>(config: {
readonly op: Op;
readonly dtype: DT;
readonly device: 'cpu';
readonly cpu: CPUKernelConfig;
readonly predicate?: KernelPredicate;
readonly fusable?: boolean;
}): KernelHandle<Op, DT, 'cpu'>function register_forward<Op extends OpName, DT extends DType>(op: Op, dtype: DT, device: 'webgpu', webgpuConfig: WebGPUKernelConfig, options?: RegisterKernelOptions): KernelHandle<Op, DT, 'webgpu'>Register a kernel with full configuration object.
Parameters
config{ readonly op: Op; readonly dtype: DT; readonly device: 'cpu'; readonly cpu: CPUKernelConfig; readonly predicate?: KernelPredicate; readonly fusable?: boolean; }
Returns
KernelHandle<Op, DT, 'cpu'>Examples
torch.library.register_forward({
op: 'add',
dtype: 'float32',
device: 'cpu',
cpu: {
forward: (a, b, out, n) => {
for (let i = 0; i < n; i++) out[i] = a[i] + b[i];
},
},
});