SoftmaxParam

Attribute	Type	Default Value	Description
axes	List[int]	list()	-

Restrictions

The length of axes must be less than that of DEFAULT_SVECTOR_SIZE. Otherwise, a length overflow exception is thrown.

Example

import torch
import torch_atb  

def softmax():
    input_npu = torch.randn(2, 16, 256, dtype=torch.float32).npu()
    print("input: ", input_npu)
    softmax_param = torch_atb.SoftmaxParam(axes = [0])
    softmax = torch_atb.Operation(softmax_param)

    def softmax_run():
        softmax_outputs = softmax.forward([input_npu])
        return softmax_outputs

    outputs = softmax_run()
    print("outputs: ", outputs)

if __name__ == "__main__":
    softmax()

Parent topic: OpParam