set_value

Description

Sets the elements that meet specific conditions in a tensor. For details, see Example.

Prototype

set_value(tensor, condition, value)

Parameters

tensor: a tvm.tensor for the input.
The supported data types include float16, int16, uint16, float32, int32, and uint32.
condition: a lambda expression to describe the scope of set_value, that is, the condition that the elements in the input tensor need to be modified.
value: a specific value. It can be a const, variable, or lambda expression.

Returns

wrapped_tensor: a tvm.tensor for the result tensor.

Restrictions

None

Availability

Atlas 200/300/500 Inference Product

Atlas Training Series Product

Example

Input tensor x:

x = [[1,2,3],
     [4,5,6],
     [7,8,9]]

Call the set_value API to assign 0 to the elements that meet the condition that the first dimension index is greater than 1 and second dimension index is greater than 0. A code example is as follows:

from tbe import tvm
from tbe import dsl
x = tvm.placeholder((3,3), dtype=dtype, name="x")
set_valued_tensor = dsl.set_value(x, lambda *i: tvm.all(i[0] > 1, i[1] > 0), 0)

The preceding code is used to set the values of x[2][1] and x[2][2] to 0. The output is as follows:

set_valued_tensor = [[1,2,3],
                     [4,5,6],
                     [7,0,0]]

Parent topic: To-Be-Deprecated APIs