LayerNormParam

import torch
import torch_atb  
import numbers

def layernorm():
    eps=1e-05
    batch, sentence_length, embedding_dim = 20, 5, 10
    embedding = torch.randn(batch, sentence_length, embedding_dim)
    embedding_npu = embedding.npu()
    normalized_shape = (embedding_dim,) if isinstance(embedding_dim, numbers.Integral) else tuple(embedding_dim)
    weight = torch.ones(normalized_shape, dtype=torch.float32).npu()
    bias = torch.zeros(normalized_shape, dtype=torch.float32).npu()
    print("embedding: ", embedding_npu)
    print("weight: ", weight)
    print("bias: ", bias)
    layer_norm_param = torch_atb.LayerNormParam(layer_type = torch_atb.LayerNormParam.LayerNormType.LAYER_NORM_NORM)
    layer_norm_param.norm_param.epsilon = eps
    layer_norm_param.norm_param.begin_norm_axis = len(normalized_shape) * -1
    layernorm = torch_atb.Operation(layer_norm_param)

    def layernorm_run():
        layernorm_outputs = layernorm.forward([embedding_npu, weight, bias])
        return layernorm_outputs

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

if __name__ == "__main__":
    layernorm()