jburtoft's picture
v1.0.0: NKI kernels for NewBERT (gte) encoders on trn2
25e8424
Raw History Blame
331 Bytes
[pytest]
testpaths = tests
markers =
sim: NKI CPU simulator (nki.simulate); needs nki, numpy, ml_dtypes. Not device evidence.
cpu: CPU-only check against the Hugging Face model; needs torch + transformers (downloads the model).
neuron: runs on a Trainium2 device under PyTorch Native (torch.compile(backend="neuron")).