diff --git a/tensorflow/core/grappler/optimizers/BUILD b/tensorflow/core/grappler/optimizers/BUILD
index 6dc6cee5bdd34ff04ade70c7b6641d093bbeac21..707d242c2b75e6da6ae4292aba1abc67559981d0 100644
--- a/tensorflow/core/grappler/optimizers/BUILD
+++ b/tensorflow/core/grappler/optimizers/BUILD
@@ -594,9 +594,9 @@ cc_library(
         ":debug_stripper",
         ":dependency_optimizer",
         ":function_optimizer",
-        ":generic_layout_optimizer",
         ":graph_optimizer",
         ":implementation_selector",
+        ":layout_optimizer",
         ":loop_optimizer",
         ":memory_optimizer",
         ":model_pruner",
diff --git a/tensorflow/core/grappler/optimizers/meta_optimizer.cc b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
index 7f1302d6b09f9e0b2e77610d059029c474d621f9..057e5a6f5c0ee9c10a90bfc524425e57384027cd 100644
--- a/tensorflow/core/grappler/optimizers/meta_optimizer.cc
+++ b/tensorflow/core/grappler/optimizers/meta_optimizer.cc
@@ -32,8 +32,8 @@ limitations under the License.
 #include "tensorflow/core/grappler/optimizers/debug_stripper.h"
 #include "tensorflow/core/grappler/optimizers/dependency_optimizer.h"
 #include "tensorflow/core/grappler/optimizers/function_optimizer.h"
-#include "tensorflow/core/grappler/optimizers/generic_layout_optimizer.h"
 #include "tensorflow/core/grappler/optimizers/implementation_selector.h"
+#include "tensorflow/core/grappler/optimizers/layout_optimizer.h"
 #include "tensorflow/core/grappler/optimizers/loop_optimizer.h"
 #include "tensorflow/core/grappler/optimizers/memory_optimizer.h"
 #include "tensorflow/core/grappler/optimizers/model_pruner.h"
@@ -121,7 +121,7 @@ std::unique_ptr<GraphOptimizer> MetaOptimizer::MakeNewOptimizer(
   MK_OPT("constfold", new ConstantFolding(cpu_device_));
   MK_OPT("shape", new ShapeOptimizer());
   MK_OPT("remap", new Remapper(cfg_.remapping()));
-  MK_OPT("layout", new GenericLayoutOptimizer());
+  MK_OPT("layout", new LayoutOptimizer());
   MK_OPT("auto_mixed_precision",
          new AutoMixedPrecision(cfg_.auto_mixed_precision()));
   MK_OPT("memory", new MemoryOptimizer(RewriterConfig::MANUAL));
@@ -193,7 +193,7 @@ Status MetaOptimizer::InitializeOptimizers(
         MakeUnique<DependencyOptimizer>(cfg_.dependency_optimization()));
   }
   if (cfg_.layout_optimizer() != RewriterConfig::OFF) {
-    optimizers->push_back(MakeUnique<GenericLayoutOptimizer>());
+    optimizers->push_back(MakeUnique<LayoutOptimizer>());
   }
   if (AutoMixedPrecisionEnabled(cfg_.auto_mixed_precision())) {
     optimizers->push_back(
@@ -267,7 +267,7 @@ Status MetaOptimizer::InitializeCustomGraphOptimizers(
       TF_RETURN_IF_ERROR(custom_optimizer->Init(&optimizer_config));
       optimizers->push_back(std::move(custom_optimizer));
     } else {
-      // If there are no custom optimizers with given name, try to initialize a
+      // If there are no custom optimizers with given name, try to initalize a
       // default optimizer. This way, custom configurable optimizers can be
       // mixed with default optimizers in any order.
       auto optimizer = MakeNewOptimizer(optimizer_config.name());
diff --git a/tensorflow/python/BUILD b/tensorflow/python/BUILD
index 93f43ab338a67dcd3198510df1501868c438aba7..c5fc23c6fccffe170251415e8c69b9bf4eab3611 100644
--- a/tensorflow/python/BUILD
+++ b/tensorflow/python/BUILD
@@ -6455,39 +6455,37 @@ cuda_py_test(
     xla_enable_strict_auto_jit = True,
 )
 
-# TODO(b/131764887) Remove once LayoutOptimizer is swapped out with GenericLayoutOptimizer.
-#
-# cuda_py_test(
-#     name = "layout_optimizer_test",
-#     size = "medium",
-#     srcs = [
-#         "grappler/layout_optimizer_test.py",
-#     ],
-#     additional_deps = [
-#         ":client_testlib",
-#         ":framework_for_generated_wrappers",
-#         ":array_ops",
-#         ":constant_op",
-#         ":dtypes",
-#         ":functional_ops",
-#         ":math_ops",
-#         ":nn",
-#         ":ops",
-#         ":random_ops",
-#         ":state_ops",
-#         ":tf_cluster",
-#         ":tf_optimizer",
-#         ":training",
-#         "//third_party/py/numpy",
-#         "//tensorflow/core:protos_all_py",
-#     ],
-#     shard_count = 10,
-#     tags = [
-#         "grappler",
-#     ],
-#     # This test will not run on XLA because it primarily tests the TF Classic flow.
-#     xla_enable_strict_auto_jit = False,
-# )
+cuda_py_test(
+    name = "layout_optimizer_test",
+    size = "medium",
+    srcs = [
+        "grappler/layout_optimizer_test.py",
+    ],
+    additional_deps = [
+        ":client_testlib",
+        ":framework_for_generated_wrappers",
+        ":array_ops",
+        ":constant_op",
+        ":dtypes",
+        ":functional_ops",
+        ":math_ops",
+        ":nn",
+        ":ops",
+        ":random_ops",
+        ":state_ops",
+        ":tf_cluster",
+        ":tf_optimizer",
+        ":training",
+        "//third_party/py/numpy",
+        "//tensorflow/core:protos_all_py",
+    ],
+    shard_count = 10,
+    tags = [
+        "grappler",
+    ],
+    # This test will not run on XLA because it primarily tests the TF Classic flow.
+    xla_enable_strict_auto_jit = False,
+)
 
 py_library(
     name = "cost_analyzer",