torch.register_scalar_forward
function register_scalar_forward<DT extends DType>(config: {
readonly op: string;
readonly dtype: DT;
readonly device: 'cpu';
readonly cpu: ScalarCPUKernelConfig;
readonly predicate?: KernelPredicate;
}): ScalarKernelHandle<string, DT, 'cpu'>function register_scalar_forward<DT extends DType>(config: {
readonly op: string;
readonly dtype: DT;
readonly device: 'webgpu';
readonly webgpu: ScalarWebGPUKernelConfig;
readonly predicate?: KernelPredicate;
}): ScalarKernelHandle<string, DT, 'webgpu'>Register a scalar operation kernel (tensor-scalar operations).
Scalar operations differ from binary operations in that one operand is always a JavaScript number, not a tensor. This enables optimized implementations that don't need to read from two tensor buffers.
Parameters
config{ readonly op: string; readonly dtype: DT; readonly device: 'cpu'; readonly cpu: ScalarCPUKernelConfig; readonly predicate?: KernelPredicate; }
Returns
ScalarKernelHandle<string, DT, 'cpu'>Examples
// Register CPU scalar kernel
torch.library.register_scalar_forward({
op: 'add_scalar',
dtype: 'float32',
device: 'cpu',
cpu: {
forward: (input, scalar, out, numel) => {
const inArr = input as Float32Array;
const outArr = out as Float32Array;
for (let i = 0; i < numel; i++) {
outArr[i] = inArr[i] + scalar;
}
},
},
});
// Register WebGPU scalar kernel
torch.library.register_scalar_forward({
op: 'add_scalar',
dtype: 'float32',
device: 'webgpu',
webgpu: {
shader: SCALAR_F32_SHADER,
entryPoint: 'add_scalar',
},
});