Check dim order in the optimized layer_norm as the portable one does

Author: SuryanshSS1011

kernels/optimized/cpu/op_native_layer_norm.cpp

@@ -149,6 +149,17 @@ std::tuple<Tensor&, Tensor&, Tensor&> opt_native_layer_norm_out(
       InvalidArgument,
       ret_val);
 
+  // Only support default dim order for now.
+  // TODO: Support other dim orders.
+  ET_KERNEL_CHECK(
+      ctx, tensor_is_default_dim_order(input), InvalidArgument, ret_val);
+
+  ET_KERNEL_CHECK(
+      ctx,
+      tensors_have_same_dim_order(input, out, mean_out, rstd_out),
+      InvalidArgument,
+      ret_val);
+
   Tensor::SizesType mean_rstd_sizes[kTensorDimensionLimit];
   size_t mean_rstd_ndim = 0;
   get_layer_norm_out_target_size(

kernels/test/op_native_layer_norm_test.cpp

@@ -452,3 +452,29 @@ TEST_F(OpNativeLayerNormTest, DynamicShapeUnbound) {
   test_dynamic_shape(
       {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND);
 }
+
+TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderDies) {
+  TensorFactory<ScalarType::Float> tf;
+
+  // mean and rstd share the input's rank with the normalized dims set to 1.
+  // All four are channels-last so the same-dim-order check passes and only the
+  // default dim order check can reject.
+  Tensor input = tf.channels_last_like(
+      tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12}));
+  Tensor out0 = tf.zeros_channels_last({1, 3, 2, 2});
+  Tensor out1 = tf.zeros_channels_last({1, 3, 2, 1});
+  Tensor out2 = tf.zeros_channels_last({1, 3, 2, 1});
+  const std::vector<int64_t> normalized_shape = {2};
+
+  ET_EXPECT_KERNEL_FAILURE(
+      context_,
+      op_native_layer_norm_out(
+          input,
+          IntArrayRef(normalized_shape.data(), normalized_shape.size()),
+          exec_aten::optional<Tensor>(),
+          exec_aten::optional<Tensor>(),
+          1e-5,
+          out0,
+          out1,
+          out2));
+}