// RUN: tf-opt -pass-pipeline='builtin.module(func.func(infeed-ops-xla-adjust-layout))' %s | FILECHECK_OPTS="" FileCheck %s

func.func @infeed_dequeue_tuple() -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>) {
  // CHECK: [[TOKEN:%.*]] = mhlo.create_token : !mhlo.token
  %0 = "mhlo.create_token"() : () -> !mhlo.token

  // CHECK:               [[INFEED:%.*]]:3 = "mhlo.infeed"([[TOKEN]]) <{
  // CHECK-SAME{LITERAL}:   infeed_config = "", layout = [[1, 3, 2, 0], [1, 2, 0]]

  ];

  let constructor = "mlir::mhlo::CreateTFXLADeviceSpecificTransformsPass()";
}

def InfeedOpsXlaAdjustLayout : Pass<"infeed-ops-xla-adjust-layout",
                                    "mlir::func::FuncOp"> {
  let summary = "Adjusts Infeed ops layout for XLA.";

  let description = [{"Adjust layouts so infeed send & receive use the same "
                    "format."}];

          if (retryMs != -1L) {
            callback.onRetryChange(retryMs)
          }
        }

        -1 -> {
          val lineEnd = source.indexOfElement(CRLF)
          if (lineEnd != -1L) {
            // Skip the line and newline
            source.skip(lineEnd)
            source.select(options)
          } else {
            return false // No more newlines.
          }
        }

std::unique_ptr<OperationPass<func::FuncOp>>
CreateTFXLADeviceSpecificTransformsPass(
    std::optional<StringRef> tf2xla_fallback_device_type = std::nullopt);

// Adjusts XLA layout for Infeed ops.
std::unique_ptr<OperationPass<func::FuncOp>>
CreateInfeedsOpsXlaAdjustLayoutPass();

#define GEN_PASS_REGISTRATION
#define GEN_PASS_DECL_INFEEDSOPSXLAADJUSTLAYOUT
#define GEN_PASS_DECL_LEGALIZETF

#include "xla/translate/mhlo_to_hlo/type_to_shape.h"

namespace mlir {

static FailureOr<std::vector<int64_t>> GetTPUInfeedLayoutFromAPI(
    RankedTensorType t) {
  // Call the TPU API to determine the right infeed layout. Note that
  // this can fail if we're not running on a TPU-enabled node.
  // TODO(kramm): Move this into a separate pass. See b/184944903
  xla::Shape old_shape = xla::TypeToShape(t);
  XLA_Shape old_shape_c = {};