This document is relevant for: Trn1, Trn2, Trn3

nki.language.where#

nki.language.where(condition, x, y, dtype=None)[source]#

Return elements chosen from x or y depending on condition.

((Similar to numpy.where))

Warning

This API is experimental and may change in future releases.

Parameters:

condition – condition tile with float values (1.0 for True, 0.0 for False).
x – tensor from which to take elements where condition is True.
y – tensor from which to take elements where condition is False.
dtype – (optional) data type to cast the output type to (see Supported Data Types for more information); if not specified, it will default to be the same as the data type of the input tile.

Returns:

tensor with elements from x or y based on condition.

Examples:

import nki.language as nl

# nki.language.where -- select 10.0 where condition is 1, else 0.0
cond = nl.full((128, 512), 1.0, dtype=nl.float32,
               buffer=nl.sbuf)
x = nl.full((128, 512), 10.0, dtype=nl.float32,
            buffer=nl.sbuf)
y = nl.full((128, 512), 0.0, dtype=nl.float32,
            buffer=nl.sbuf)
result = nl.where(cond, x, y)
expected = nl.full((128, 512), 10.0, dtype=nl.float32,
                   buffer=nl.sbuf)
assert nl.equal(result, expected)

# nki.language.where -- select 5.0 where condition is 0
cond = nl.full((128, 512), 0.0, dtype=nl.float32,
               buffer=nl.sbuf)
x = nl.full((128, 512), 10.0, dtype=nl.float32,
            buffer=nl.sbuf)
y = nl.full((128, 512), 5.0, dtype=nl.float32,
            buffer=nl.sbuf)
result = nl.where(cond, x, y)
expected = nl.full((128, 512), 5.0, dtype=nl.float32,
                   buffer=nl.sbuf)
assert nl.equal(result, expected)

This document is relevant for: Trn1, Trn2, Trn3

nki.language.where

Contents

nki.language.where#