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