|
|
|
@@ -37,6 +37,8 @@ int LayerNormInferShape(const TensorC *const *inputs, size_t inputs_size, Tensor |
|
|
|
if (!param->op_parameter_.infer_flag_) { |
|
|
|
return NNACL_INFER_INVALID; |
|
|
|
} |
|
|
|
param->begin_norm_axis_ = |
|
|
|
param->begin_norm_axis_ < 0 ? param->begin_norm_axis_ + input->shape_size_ : param->begin_norm_axis_; |
|
|
|
SetShapeTensor(output, input); |
|
|
|
// take care of other outputs |
|
|
|
if (outputs_size == 3) { |
|
|
|
@@ -45,10 +47,9 @@ int LayerNormInferShape(const TensorC *const *inputs, size_t inputs_size, Tensor |
|
|
|
SetDataTypeFormat(output_mean, input); |
|
|
|
SetDataTypeFormat(output_var, input); |
|
|
|
int size = 0; |
|
|
|
for (int i = param->begin_norm_axis_; i < input->shape_size_; i++) { |
|
|
|
output_mean->shape_[size] = input->shape_[i]; |
|
|
|
output_var->shape_[size] = input->shape_[i]; |
|
|
|
size++; |
|
|
|
for (; size < param->begin_norm_axis_; size++) { |
|
|
|
output_mean->shape_[size] = input->shape_[size]; |
|
|
|
output_var->shape_[size] = input->shape_[size]; |
|
|
|
} |
|
|
|
output_mean->shape_size_ = size; |
|
|
|
output_var->shape_size_ = size; |
|
|
|
|