[KernelGen][JIT] Enable JIT-compiler OnesLike kernels
PiperOrigin-RevId: 412088106
Change-Id: I761a68252c24aa6305a51f8c3885dc7d4d002909
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_op_ones_like.cc b/tensorflow/core/kernels/mlir_generated/gpu_op_ones_like.cc
index c79c4e6..9121dcd 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_op_ones_like.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_op_ones_like.cc
@@ -25,13 +25,11 @@
GENERATE_AND_REGISTER_UNARY_GPU_KERNEL(OnesLike, DT_INT64);
// These kernels are JIT-compiled.
-#if defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(OnesLike, DT_INT8);
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(OnesLike, DT_INT16);
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(OnesLike, DT_UINT8);
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(OnesLike, DT_UINT16);
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(OnesLike, DT_UINT32);
GENERATE_AND_REGISTER_UNARY_JIT_GPU_KERNEL(OnesLike, DT_UINT64);
-#endif
} // namespace tensorflow
diff --git a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
index d139449..773ffb8 100644
--- a/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
+++ b/tensorflow/core/kernels/mlir_generated/gpu_unary_ops_test.cc
@@ -784,8 +784,7 @@
test::OpsTestConfig().ExpectStrictlyEqual())
// These kernels are JIT-compiled.
-#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED) && \
- defined(MLIR_GENERATED_EXPERIMENTAL_KERNELS_ENABLED)
+#if defined(MLIR_GENERATED_GPU_KERNELS_ENABLED)
GENERATE_DEFAULT_TEST(OnesLike, DT_INT8, DT_INT8, baseline_ones_like,
test::OpsTestConfig().ExpectStrictlyEqual())
GENERATE_DEFAULT_TEST(OnesLike, DT_INT16, DT_INT16, baseline_ones_like,
