torch.register_device
function register_device<const Name extends string>(config: DeviceConfig<Name>): DeviceHandle<Name>Register a custom device backend.
Allows third-party packages to add new device types like WebGL, WASM, or native backends. The device becomes available for kernel registration and tensor operations.
Parameters
configDeviceConfig<Name>- Device configuration
Returns
DeviceHandle<Name>– Handle to the registered deviceExamples
// First, extend the DeviceRegistry type:
declare module '@torchjsorg/torch.js' {
interface DeviceRegistry {
webgl: 'webgl';
}
}
// Then register the device:
torch.library.register_device({
name: 'webgl',
initialize: async () => {
const canvas = document.createElement('canvas');
const gl = canvas.getContext('webgl2');
if (!gl) throw new Error('WebGL2 not available');
return {
name: 'webgl',
ready: true,
context: gl,
destroy: () => canvas.remove(),
};
},
capabilities: {
compute: false,
storage: true,
maxBufferSize: 256 * 1024 * 1024,
maxWorkgroupSize: [0, 0, 0],
maxWorkgroupsPerDimension: 0,
supportedDTypes: ['float32'],
},
createBuffer: (size, usage) => { ... },
dispatch: (config) => { ... },
sync: async () => { ... },
});
// Now you can register kernels for the device:
torch.library.register_forward({
op: 'add',
dtype: 'float32',
device: 'webgl',
webgl: { program: addProgram, uniforms: {} },
});