From 9e74ca45c97125cd7afd03858fc845e4d199a5c8 Mon Sep 17 00:00:00 2001 From: George Necula Date: Mon, 26 Jun 2023 22:02:03 -0700 Subject: [PATCH] [shape_poly] Shape polymorphism support for approx_top_k PiperOrigin-RevId: 543633818 --- tensorflow/compiler/xla/python/xla_client.py | 2 +- .../compiler/xla/translate/mhlo_to_hlo/mlir_hlo_to_hlo.cc | 4 +++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorflow/compiler/xla/python/xla_client.py b/tensorflow/compiler/xla/python/xla_client.py index f7fb7921ceba4d..61fd9624e7f423 100644 --- 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 index 1070265dc0f22a..480ff12fefae04 100644 --- 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))