// 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]]

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 = {};

  }
};

// Converts InfeedDequeueTuple to XLA HLO create_token, infeed and
// get_tuple_element ops.
//
// All HLO infeed ops expect a HLO token type operand and produce a tuple
// containing a token. This HLO token type is used to order multiple infeed
// operations within a computation. The token type can come from other
// infeed/outfeed/send/recv ops or can be generated using create_token op with

    TF_Tensor:$output
  );

  TF_DerivedOperandTypeAttr T = TF_DerivedOperandTypeAttr<0>;
}

def TF_InfeedDequeueTupleOp : TF_Op<"InfeedDequeueTuple", []> {
  let summary = "Fetches multiple values from infeed as an XLA tuple.";

  let arguments = (ins
    OptionalAttr<StrAttr>:$_XlaSharding,
    OptionalAttr<ArrayAttr>:$layouts
  );

  let results = (outs
    Variadic<TF_Tensor>:$outputs
  );

// CHECK: [[TOKEN:%.*]] = mhlo.create_token  : !mhlo.token
// CHECK: [[INFEED:%.*]]:3 = "mhlo.infeed"([[TOKEN]]) <{infeed_config = ""{{.*}}}> : (!mhlo.token) -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>, !mhlo.token)
// CHECK: return [[INFEED]]#0, [[INFEED]]#1
  %0:2 = "tf.InfeedDequeueTuple"() : () -> (tensor<1x8x4x4xi32>, tensor<1x100x1xf32>)

  }];

  let arguments = (ins
    TF_ShapeAttr:$shape
  );

  let results = (outs
    Res<TF_Tensor, [{A tensor that will be provided using the infeed mechanism.}]>:$output
  );

  TF_DerivedResultTypeAttr dtype = TF_DerivedResultTypeAttr<0>;
}

def TF_InitializeTableOp : TF_Op<"InitializeTable", []> {
  let summary = [{