See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

// This file implements logic for legalizing HLO to TensorFlow.

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <cstdlib>
#include <functional>
#include <memory>
#include <numeric>
#include <optional>
#include <string>

        input_output_alias);
    OP_REQUIRES_OK(ctx, execution_inputs.status());
  }

  xla::ExecutableRunOptions run_options;

  // Host callbacks used for HLO send/recv.
  xla::SendDeviceMemoryFunction send_function =
      GetSendDeviceMemoryFunction(ctx, key);
  run_options.set_send_device_memory_function(&send_function);
  xla::RecvDeviceMemoryFunction recv_function =

          : "";

  const string include_hlo_profile_printer_data_proto =
      opts.gen_hlo_profile_printer_data
          ? R"(#include "xla/service/hlo_profile_printer_data.pb.h")"
          : "";

  // When HLO profiling is disabled we only forward declare the
  // HloProfilePrinter protobuf.  So we can only conditionally emit this code
  // calling HloProfilePrinter::profile_counters_size.
  const string assign_profile_counters_size =

        "@com_google_absl//absl/status",
        "@com_google_absl//absl/types:span",
        "@local_tsl//tsl/framework:device_id_utils",
        "@local_xla//xla:executable_run_options",
        "@local_xla//xla/hlo/ir:hlo",
        "@local_xla//xla/pjrt:pjrt_client",
        "@local_xla//xla/pjrt:tf_pjrt_client",
        "@local_xla//xla/service:compiler",
        "@local_xla//xla/service:executable",

  Type new_arg_type = tensorflow::GetTypeFromTFTensorShape(shape, element_type);

  if (auto input_ty = mlir::dyn_cast<RankedTensorType>(old_arg_type)) {
    ArrayRef<int64_t> bounds = hlo::encodingToBounds(input_ty.getEncoding());
    // The input type has bounded dynamic dimension.
    if (!bounds.empty()) {
      SmallVector<int64_t> new_bounds(bounds.begin(), bounds.end());

  let description = [{
    This pass looks for the usage of the result of TPUCopyWithDynamicShapeOp
    and sets the shape of these inputs to be dynamic shaped. This will ensure
    that the generated HLO program is correctly reflecting the dynamic shape.
  }];

  // Required for mhlo bounded shape extension.
  let dependentDialects = ["mhlo::MhloDialect"];

  func.return %0 : tensor<5x11xf32>
}

// Verify that MatMul with ranked inputs are lowered to HLO.
// CHECK-LABEL: matmul_ranked
func.func @matmul_ranked(%a: tensor<?x7xf32>, %b: tensor<7x?xf32>) -> tensor<?x?xf32> {
  // CHECK: "mhlo.dot"

### XLA Support:

*   xla.experimental.compile is deprecated, use
    `tf.function(experimental_compile=True)` instead.
*   Adds `tf.function.experimental_get_compiler_ir` which returns compiler IR
    (currently 'hlo' and 'optimized_hlo') for given input for given function.

### Security:

*   Fixes an undefined behavior causing a segfault in `tf.raw_ops.Switch`,

// NOTE: Assertions have been autogenerated by utils/generate-test-checks.py
// RUN: odml-to-stablehlo-opt -tf-legalize-hlo -verify-diagnostics %s | FileCheck %s

// CHECK-LABEL:   func @biasAdd_NHWC(
// CHECK-SAME:                       %[[VAL_0:.*]]: tensor<1x32x10x32xi32>,
// CHECK-SAME:                       %[[VAL_1:.*]]: tensor<32xi32>) -> tensor<1x32x10x32xi32> {

  TF_DerivedResultTypeAttr preferred_element_type = TF_DerivedResultTypeAttr<0>;

  let hasVerifier = 1;
}

def TF_XlaCustomCallV2Op : TF_Op<"XlaCustomCallV2", [Pure]> {
  let summary = "Emits an HLO `CustomCall` operation with multiple outputs.";

  let description = [{
As opposed to `XlaCustomCall`, this operation supports multiple outputs.

See `CustomCall` specification at