torch.mps.load_metallib

torch.mps.load_metallib(source)

Loads a precompiled Metal library (.metallib) and returns a shader library object that allows invoking kernels defined in it.

Parameters:

source – Either raw metallib bytes (bytes/bytearray) or a filesystem path (str/os.PathLike) to a .metallib file.

This is useful for loading Metal libraries compiled ahead of time or generated by external tools (e.g. Triton, MetalASM).

Example:

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