[shape_poly] Shape polymorphism support for approx_top_k

PiperOrigin-RevId: 543633818
tensorflow · Jun 27, 2023 · 9e74ca4 · 9e74ca4
1 parent ab87e62
commit 9e74ca4
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py
@@ -47,7 +47,7 @@
 _version = 165
 
 # Version number for MLIR:Python components.
-mlir_api_version = 50
+mlir_api_version = 51
 
 xla_platform_names = {
     'cpu': 'Host',

diff --git a/tensorflow/compiler/xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc b/tensorflow/compiler/xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc
@@ -1675,6 +1675,7 @@ LogicalResult ExportXlaOp(CustomCallOp op, OpLoweringContext ctx) {
   // The attributes supported by the ApproxTopK custom_call are:
   //
   //  - called_computation : This indicates the comparator for scoring entries
+  //  - has_side_effect: always False
   //  - api_version : always 4, the typed FFI API
   //  - backend_config : The actual arguments to ApproxTopK. This includes
   //    + top_k:i64 : the number of results to return
@@ -1738,7 +1739,8 @@ LogicalResult ExportXlaOp(CustomCallOp op, OpLoweringContext ctx) {
     auto isSupportedAttrName = [](NamedAttribute attr) {
       auto name = attr.getName();
       return name == "call_target_name" || name == "backend_config" ||
-             name == "api_version" || name == "called_computations";
+             name == "api_version" || name == "called_computations" ||
+             name == "has_side_effect";
     };
     for (const auto& attr : op->getAttrs()) {
       if (!isSupportedAttrName(attr))