torch.mps.load_metallib

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

Загружает предварительно скомпилированную библиотеку Metal (.metallib) и возвращает объект библиотеки шейдеров, позволяющий вызывать определённые в ней ядра.

Параметры:

source – Сырые байты metallib (bytes/bytearray) или путь в файловой системе (str/os.PathLike) к файлу .metallib.

Это удобно для загрузки библиотек Metal, скомпилированных заранее или созданных внешними инструментами (например, Triton, MetalASM).

Пример:

>>> lib = torch.mps.load_metallib("kernels.metallib")
>>> x = torch.ones(16, device="mps")
>>> lib.square(x)