This document is relevant for: Trn2, Trn3
nki.language.shared_constant#
- nki.language.shared_constant(constant)[source]#
Create a tensor in shared HBM initialized with constant data.
The constant is embedded in the compiled binary and loaded to HBM at model load time. With LNC=2, both cores share the same constant; the data must not diverge across cores.
Supported element types: float32, float16, bfloat16, int32, int16, int8, uint32, uint16, uint8, float8_e4m3fn, float8_e5m2, float8_e8m0fnu. Packed types (float8_e4m3fn_x4, float8_e5m2_x4, float4_e2m1fn_x4) and tfloat32 are supported at the MLIR level but not yet tested end-to-end on hardware.
- Parameters:
constant – the constant data. Can be a numpy array or a file path to a
.npyfile.- Returns:
an NkiTensor in shared_hbm containing the constant data.
This document is relevant for: Trn2, Trn3