Cast Compute APIs
Cast compute APIs are used to round the input tensor element-wise based on certain rules.
ceil
Description
Rounds up a raw_tensor element-wise.
The API is defined in python/site-packages/te/lang/cce/te_compute/cast_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16
Prototype
te.lang.cce.ceil(raw_tensor)
Parameters
raw_tensor: a tvm.tensor for the input
Returns
res_tensor: a tvm.tensor of type int32 for the result
Example
shape = (1024,1024) input_dtype = "float16" data = tvm.placeholder(shape, name="data", dtype=input_dtype) res = te.lang.cce.ceil(data)
floor
Description
Rounds down a raw_tensor element-wise.
The API is defined in python/site-packages/te/lang/cce/te_compute/cast_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16
Prototype
te.lang.cce.floor(raw_tensor)
Parameters
raw_tensor: a tvm.tensor for the input tensor
Returns
res_tensor: a tvm.tensor for the result tensor of type int32
Example
shape = (1024,1024) input_dtype = "float16" data = tvm.placeholder(shape, name="data", dtype=input_dtype) res = te.lang.cce.floor(data)
round
Description
Performs banker's rounding on a raw_tensor element-wise. For example, the constant interval between adjacent members is 0.5, then 1.5 rounds up to 2, and 2.5 rounds down to 2.0.
The API is defined in python/site-packages/te/lang/cce/te_compute/cast_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor: float16
Prototype
te.lang.cce.round(raw_tensor)
Parameters
raw_tensor: a tvm.tensor for the input tensor
Returns
res_tensor: a tvm.tensor for the result tensor of type int32
Example
shape = (1024,1024) input_dtype = "float16" data = tvm.placeholder(shape, name="data", dtype=input_dtype) res = te.lang.cce.round(data)
trunc
Description
Rounds a raw_tensor towards 0 element-wise. For example, –1.9 rounds up to –1, and 1.9 rounds down to 1.
The API is defined in python/site-packages/te/lang/cce/te_compute/cast_compute.py in the ATC installation path.
Restrictions
Ascend 310 AI Processor does not support this API.
Prototype
te.lang.cce.trunc(raw_tensor)
Parameters
raw_tensor: a tvm.tensor for the input tensor
Returns
res_tensor: a tvm.tensor for the result tensor of type int32
Example
shape = (1024,1024) input_dtype = "float16" data = tvm.placeholder(shape, name="data", dtype=input_dtype) res = te.lang.cce.trunc(data)
