nki.language.shared_constant#

nki.language.shared_constant(constant, dtype=None)[source]#

Create a tensor filled with compile-time constant data.

This function creates a tensor that contains constant data specified by a trace-time tensor. The constant data is evaluated at compile time and shared across all instances where the same constant is used, making it memory efficient for frequently used constant values.

Parameters:

  • constant (nki.tensor.TraceTimeTensor) – A trace-time tensor containing the constant data to be filled into the output tensor. This can be created using functions from nki.tensor such as nki.tensor.zeros(), nki.tensor.identity(), or nki.tensor.arange().

  • dtype (nki.language dtype) – The data type of the output tensor. Must be specified. Only types that can be serialized to npy files are supported. See Supported Data Types for supported data types.

Returns:

A tensor containing the constant data with the specified dtype.

Return type:

Tensor

Note

The constant tensor is shared across all uses of the same constant data and dtype, which helps reduce memory usage in the compiled kernel.

Examples:

Create a constant identity matrix:

import nki.tensor as ntensor
import nki.language as nl

# Create a 128x128 identity matrix as a shared constant
identity_matrix = nl.shared_constant(
    ntensor.identity(128, dtype=nl.int8),
    dtype=nl.float16
)

Create a constant tensor with sequential values:

# Create a constant tensor with values [0, 1, 2, ..., 31]
sequential_values = nl.shared_constant(
    ntensor.arange(0, 32, 1, dtype=nl.int32),
    dtype=nl.float32
)

Create a constant tensor with arithmetic operations:

# Create a constant tensor filled with ones
ones_tensor = nl.shared_constant(
    ntensor.zeros((64, 64), dtype=nl.int8) + 1,
    dtype=nl.int16
)