nki.compiler.psum.auto_alloc — AWS Neuron Documentation (original) (raw)

This document is relevant for: Inf2, Trn1, Trn2

nki.compiler.psum.auto_alloc#

nki.compiler.psum.auto_alloc()[source]#

Returns a maker to indicate the tensor should be automatically allocated by compiler. All PSUM tensors in a kernel must either all be marked as auto_alloc(), or all be allocated with alloc or mod_alloc.

Initialize a tensor with buffer=nl.psum is equivalent to buffer=ncc.psum.auto_alloc().

This document is relevant for: Inf2, Trn1, Trn2