nki — AWS Neuron Documentation (original) (raw)
Contents
This document is relevant for: Inf2, Trn1, Trn2
nki#
Decorators#
| jit | This decorator compiles a function to run on NeuronDevices. |
|---|---|
| benchmark | Benchmark a NKI kernel on a NeuronDevice by using nki.benchmark as a decorator. |
| profile | Profile a NKI kernel on a NeuronDevice by using nki.profile as a decorator. |
| baremetal | Compile and run a NKI kernel on NeuronDevice without involving ML frameworks such as PyTorch and JAX. |
| simulate_kernel | Simulate a nki kernel on CPU using a built-in simulator in Neuron Compiler. |
Types#
| tensor | A tensor object represents a multidimensional, homogeneous array of fixed-size items |
|---|
This document is relevant for: Inf2, Trn1, Trn2