torch.mps.compile_shader

torch.mps.compile_shader(source) [исходный код]

Компилирует вычислительный шейдер из исходного кода и позволяет вызывать определённые в нём ядра прямо из среды выполнения Python. Пример:

>>> lib = torch.mps.compile_shader(
... "kernel void full(device float* out, constant float& val, uint idx [[thread_position_in_grid]]) { out[idx] = val; }"
...  )
>>> x = torch.zeros(16, device="mps")
>>> lib.full(x, 3.14)