diff --git a/include/tim/vx/ops/LayerNormalization.json b/include/tim/vx/ops/LayerNormalization.json
index 53d685b..7dbd359 100755
--- a/include/tim/vx/ops/LayerNormalization.json
+++ b/include/tim/vx/ops/LayerNormalization.json
@@ -3,9 +3,7 @@
         "parameters":
             [
                 {"name": "axis",
-                 "dtype": "int32_t",
-                 "Optional": "true",
-                 "default": "0"
+                 "dtype": "int32_t"
                 },
                 {"name": "eps",
                  "dtype": "float",
diff --git a/include/tim/vx/ops/layernormalization.h b/include/tim/vx/ops/layernormalization.h
index 0cdc46b..9dceb93 100644
--- a/include/tim/vx/ops/layernormalization.h
+++ b/include/tim/vx/ops/layernormalization.h
@@ -32,7 +32,7 @@ namespace vx {
 namespace ops {
 class LayerNormalization : public BuiltinOp {
  public:
-  LayerNormalization(Graph* graph, int32_t axis = 0, float eps = 1e-5f);
+  LayerNormalization(Graph* graph, int32_t axis, float eps = 1e-5f);
 
   std::shared_ptr<Operation> Clone(std::shared_ptr<Graph>& graph) const override;
 
diff --git a/src/tim/vx/ops/layernormalization.cc b/src/tim/vx/ops/layernormalization.cc
index 038a2c0..9976405 100644
--- a/src/tim/vx/ops/layernormalization.cc
+++ b/src/tim/vx/ops/layernormalization.cc
@@ -32,11 +32,7 @@ namespace vx {
 namespace ops {
 LayerNormalization::LayerNormalization(Graph* graph, int32_t axis, float eps)
     : BuiltinOp(graph, VSI_NN_OP_LAYER_NORM), axis_(axis), eps_(eps) {
-  // Layer normalization shares the parameters of instance normalization.
-  if (axis != 0) {
-    VSILOGE("Layer norm only support axis 0.");
-    assert(false);
-  }
+  this->impl()->node()->nn_param.layernorm.axis = axis_;
   this->impl()->node()->nn_param.layernorm.eps = eps_;
 }
 
diff --git a/src/tim/vx/ops/layernormalization_test.cc b/src/tim/vx/ops/layernormalization_test.cc
index 3c7c9df..9b509ce 100644
--- a/src/tim/vx/ops/layernormalization_test.cc
+++ b/src/tim/vx/ops/layernormalization_test.cc
@@ -142,6 +142,61 @@ TEST(LayerNorm, axis_0_shape_2_3_6_1_float) {
     EXPECT_TRUE(ArraysMatch(golden, output, 1e-5f));
 }
 
+TEST(LayerNorm, axis_2_shape_4_2_3_1_float) {
+    auto ctx = tim::vx::Context::Create();
+    auto graph = ctx->CreateGraph();
+
+    tim::vx::ShapeType io_shape({4, 2, 3, 1});
+    tim::vx::ShapeType param_shape({1,1,3,1});
+    tim::vx::TensorSpec input_spec(tim::vx::DataType::FLOAT32,
+                            io_shape, tim::vx::TensorAttribute::INPUT);
+    tim::vx::TensorSpec param_spec(tim::vx::DataType::FLOAT32,
+                            param_shape, tim::vx::TensorAttribute::INPUT);
+    tim::vx::TensorSpec output_spec(tim::vx::DataType::FLOAT32,
+                            io_shape, tim::vx::TensorAttribute::OUTPUT);
+
+    auto input_tensor = graph->CreateTensor(input_spec);
+    auto gamma_tensor = graph->CreateTensor(param_spec);
+    auto beta_tensor = graph->CreateTensor(param_spec);
+    auto output_tensor = graph->CreateTensor(output_spec);
+
+    std::vector<float> in_data = {
+        1, 2, 3, 4,
+        5, 6, 7, 8,
+        9, 10, 11, 12,
+        13, 14, 15, 16,
+        17, 18, 19, 20,
+        21, 22, 23, 24};
+    std::vector<float> gamma = {
+        1.0f, 1.0f, 1.0f
+    };
+    std::vector<float> beta = {
+        .0f, .0f, .0f
+    };
+    std::vector<float> golden = {
+        -1.22473, -1.22473, -1.22473, -1.22473,
+        -1.22473, -1.22473, -1.22473, -1.22473,
+         0, 0, 0, 0,
+         0, 0, 0, 0,
+         1.22473, 1.22473, 1.22473, 1.22473,
+         1.22473, 1.22473, 1.22473, 1.22473
+    };
+
+    EXPECT_TRUE(input_tensor->CopyDataToTensor(in_data.data(), in_data.size() * sizeof(float)));
+    EXPECT_TRUE(gamma_tensor->CopyDataToTensor(gamma.data(), gamma.size() * sizeof(float)));
+    EXPECT_TRUE(beta_tensor->CopyDataToTensor(beta.data(), beta.size() * sizeof(float)));
+
+    auto op = graph->CreateOperation<tim::vx::ops::LayerNormalization>(2, 0.001);
+    (*op).BindInputs({input_tensor, beta_tensor, gamma_tensor}).BindOutputs({output_tensor});
+
+    EXPECT_TRUE(graph->Compile());
+    EXPECT_TRUE(graph->Run());
+
+    std::vector<float> output(24);
+    EXPECT_TRUE(output_tensor->CopyDataFromTensor(output.data()));
+    EXPECT_TRUE(ArraysMatch(golden, output, 1e-5f));
+}
+
 #if 0
 // Fail case
 TEST(LayerNorm, axis_0_shape_3_6_1_uint8) {