Operation
In the scenario where a user calls a single OP, the Operation provides APIs for creating, executing, and the OP. You can directly perform operations related to the OP through this object.
Creating an Operation
API |
CLASS torch_atb.Operation(OpParam) |
|---|---|
API Description |
Creates a single OP. |
(Optional) Input Requirement/Parameter |
OpParam |
(Optional) Output Requirement/Parameter |
torch_atb.Operation |
Currently, the following input parameter types are supported: LayerNormParam, ElewiseParam, LinearParam, SoftmaxParam, RopeParam, SelfAttentionParam, PagedAttentionParam, SplitParam, ReshapeAndCacheParam, GatherParam, ActivationParam, and RmsNormParam.
name
Description |
Attribute of torch_atb.Operation, which is used to obtain the name of torch_atb.Operation. |
|---|---|
Prototype |
torch_atb.Operation.name -> str |
Parameter |
NA |
Returns |
Name of an Operation. |
Note |
If the API fails to be called, an exception is thrown and the API stops running. |
Example:
linear = torch_atb.Operation(linear_param) print(linear.name)
input_num
Description |
Attribute of torch_atb.Operation, which is used to obtain the number of input tensors. |
|---|---|
Prototype |
torch_atb.Operation.input_num ->int |
Parameter |
NA |
Returns |
Number of input tensors. |
Note |
If the API fails to be called, an exception is thrown and the API stops running. |
Example:
linear = torch_atb.Operation(linear_param) print(linear.input_num)
output_num
Description |
Attribute of torch_atb.Operation, which is used to obtain the number of output tensors. |
|---|---|
Prototype |
torch_atb.Operation.output_num ->int |
Parameter |
NA |
Returns |
Number of output tensors. |
Note |
If the API fails to be called, an exception is thrown and the API stops running. |
Example:
linear = torch_atb.Operation(linear_param) print(linear.output_num)
forward
Description |
Method of torch_atb.Operation, which is used to execute an Operation. |
|---|---|
Prototype |
torch_atb.Operation.forward(List:torch::Tensor) -> List:torch::Tensor |
Parameter |
Input tensor of type List[torch.Tensor]. |
Returns |
Output tensor of type List[torch.Tensor]. |
Note |
If the API fails to be called, an exception is thrown and the API stops running. |
Example:
input = torch.randn(m, k, dtype=torch.float16).npu()
weight = torch.randn(k, n, dtype=torch.float16).npu()
linear_param = torch_atb.LinearParam()
linear_param.has_bias = False
linear_param.transpose_b = False
linear = torch_atb.Operation(linear_param)
linear_outputs = linear.forward([input, weight])
print("linear_outputs : ", linear_outputs )