inplace_sub

Description

Subtracts rhs into specified rows of lhs.

See the following example.

res = lhs
res[ids,:] -= rhs
return res

Prototype

inplace_sub(lhs, inplace_ids, rhs)

Parameters

• lhs: a tensor for the left input
• inplace_ids: an int or a list of ints within the range of [0, the first dimension of lhs]. The length must be the same as that of the first dimension of rhs.
• rhs: a tensor or scalar for the right input. Has the same dimension sizes as lhs except the first dimension. If inplace_ids is an int, rhs has one less dimension than lhs. For example, lhs is (10, 1024), inplace_ids is [5], and rhs is (1, 1024); lhs is (10, 1024), inplace_ids is 5, and rhs is (1024,).
• The following data types are supported:

  Atlas 200/300/500 Inference Product : supports float16, float32, and int32.

  Atlas Training Series Product : supports float16, float32, and int32.

Returns

res_tensor: result tensor.

Restrictions

• This API cannot be used in conjunction with other TBE DSL APIs.
• The maximum value of the first dimension of rhs is 7934. A value larger than 7934 cannot be processed.
• If the first dimension value of rhs is 5000 or larger, a core dump resulting from OS stack overflow may occur. In this case, you can run the ulimit -s command to increase the stack space, for example, from 8192 to 81920.

Applicability

Atlas 200/300/500 Inference Product

Atlas Training Series Product

Example

from tbe import tvm
from tbe import dsl
input_dtype = "float16"
dataA = tvm.placeholder((6,1024), name="dataA", dtype=input_dtype)
dataB = tvm.placeholder((5,1024), name="dataB", dtype=input_dtype)
inplace_ids = [1,1,4,2,2]
res = dsl.inplace_sub(dataA, inplace_ids, dataB)
res.shape = (6,1024)
# res[0] = dataA[0]
# res[1] = dataA[1] - dataB[0] - dataB[1]
# res[2] = dataA[2] - dataB[3] - dataB[4]
# res[3] = dataA[3]
# res[4] = dataA[4] - dataB[2]
# res[5] = dataA[5]