diff --git a/kernels/portable/cpu/util/normalization_ops_util.cpp b/kernels/portable/cpu/util/normalization_ops_util.cpp index 4adcf02b303..dedab427ae7 100644 --- a/kernels/portable/cpu/util/normalization_ops_util.cpp +++ b/kernels/portable/cpu/util/normalization_ops_util.cpp @@ -7,6 +7,7 @@ */ #include +#include #include #include @@ -92,6 +93,11 @@ bool check_layer_norm_args( ", ndim = %zu", in.dim(), ndim); + ET_CHECK_OR_RETURN_FALSE( + ndim <= kTensorDimensionLimit, + "Expected normalized shape to have at most %zu dimensions but it had %zu", + kTensorDimensionLimit, + ndim); size_t shift = in.dim() - ndim; for (const auto d : c10::irange(ndim)) { ET_CHECK_OR_RETURN_FALSE( @@ -103,7 +109,7 @@ bool check_layer_norm_args( d, normalized_shape[d]); } - executorch::aten::SizesType shape[ndim]; + std::array shape; for (const auto i : c10::irange(ndim)) { shape[i] = static_cast(normalized_shape[i]); } @@ -111,12 +117,12 @@ bool check_layer_norm_args( if (weight.has_value()) { ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, weight.value())); ET_LOG_AND_RETURN_IF_FALSE( - tensor_has_expected_size(weight.value(), {shape, ndim})); + tensor_has_expected_size(weight.value(), {shape.data(), ndim})); } if (bias.has_value()) { ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, bias.value())); ET_LOG_AND_RETURN_IF_FALSE( - tensor_has_expected_size(bias.value(), {shape, ndim})); + tensor_has_expected_size(bias.value(), {shape.data(), ndim})); } ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, out)); ET_LOG_AND_RETURN_IF_FALSE(tensors_have_same_dtype(in, mean_out)); diff --git a/test/build_optimized_size_test.sh b/test/build_optimized_size_test.sh index d7b055559f0..3141a29e9f1 100644 --- a/test/build_optimized_size_test.sh +++ b/test/build_optimized_size_test.sh @@ -41,12 +41,15 @@ test_cmake_size_test() { echo 'ExecuTorch with no ops binary size, unstripped:' ls -al cmake-out/test/size_test + size cmake-out/test/size_test echo 'ExecuTorch with portable ops binary size, unstripped:' ls -al cmake-out/test/size_test_all_ops + size cmake-out/test/size_test_all_ops echo 'ExecuTorch with optimized ops binary size, unstripped:' ls -al cmake-out/test/size_test_all_optimized_ops + size cmake-out/test/size_test_all_optimized_ops } if [[ -z $PYTHON_EXECUTABLE ]]; then diff --git a/test/build_size_test.sh b/test/build_size_test.sh index d020ab58c95..baeef5a849e 100644 --- a/test/build_size_test.sh +++ b/test/build_size_test.sh @@ -46,9 +46,11 @@ test_cmake_size_test() { echo 'ExecuTorch with no ops binary size, unstripped:' ls -al cmake-out/test/size_test + size cmake-out/test/size_test echo 'ExecuTorch with portable ops binary size, unstripped:' ls -al cmake-out/test/size_test_all_ops + size cmake-out/test/size_test_all_ops } if [[ -z $PYTHON_EXECUTABLE ]]; then