op->getLoc(), list_type, variadic_operands);
        new_operands.push_back(build_list_op.getOut());
      } else {
        auto attr_name = compose_func.getArgAttrOfType<StringAttr>(
            arg.index(), kAttrArgumentNameAttr);
        auto attribute = op->getAttr(attr_name.getValue());
        if (!attribute) {
          attribute =

      llvm::dyn_cast<tf_device::LaunchOp>(preprocess_op->getParentOp());
  // Device of original launch looked up before moving the preprocessing ops
  // around.
  StringAttr device =
      original_launch_op->getAttrOfType<StringAttr>(kDeviceAttr);

  if (!device) {
    return original_launch_op->emitOpError()
           << "Launch op has an invalid device attribute.";
  }

using mlir::Builder;
using mlir::IntegerAttr;
using mlir::LogicalResult;
using mlir::ModuleOp;
using mlir::Operation;
using mlir::OpOperand;
using mlir::OpResult;
using mlir::RankedTensorType;
using mlir::StringAttr;
using mlir::Value;
using mlir::WalkResult;

constexpr char kReplicateSharding[] = "";
constexpr char kShardingAttr[] = "mhlo.sharding";
constexpr char kUseSpmdAttr[] = "use_spmd_for_xla_partitioning";

namespace {

using mlir::Block;
using mlir::BlockArgument;
using mlir::BoolAttr;
using mlir::ModuleOp;
using mlir::OpBuilder;
using mlir::Operation;
using mlir::OperationPass;
using mlir::Region;
using mlir::StringAttr;
using mlir::Type;
using mlir::Value;
using mlir::WalkResult;
using mlir::func::FuncOp;

constexpr char kXlaMapOutsideCompilationAttr[] = "_xla_map_outside_compilation";

  }
  const auto empty_array_attr = ArrayAttr::get(ctx, {});
  // TODO: b/310291615 - find a better way for platform support.
  const auto platforms = ArrayAttr::get(
      ctx,
      {StringAttr::get(ctx, kPlatformCpu), StringAttr::get(ctx, kPlatformTpu)});

  auto xla_call_module_op = builder.create<TF::XlaCallModuleOp>(
      module_op.getLoc(), /*output=*/result_types,
      /*args=*/inputs,

class IsFusedOpEndsWith<string OpName> : AttrConstraint<
  CPred<"!$_self.cast<ArrayAttr>().empty() && "
        "$_self.cast<ArrayAttr>()[$_self.cast<ArrayAttr>().size() - 1]."
        "cast<::mlir::StringAttr>().str() == \"" # OpName # "\"">,
  "Matching fused '" # OpName # "' op at the end">;

//===----------------------------------------------------------------------===//
// Pattern rules for lifting ops as functions

 */
bool IsResourceMergeable(Attribute& resource_attr, Attribute& device_attr) {
  return resource_attr &&
         ((resource_attr == device_attr) ||
          (resource_attr.cast<mlir::StringAttr>().getValue().find(
               "COMPOSITE") != llvm::StringRef::npos));
}

// Finds the variable access info for a TPUExecute op.
//  - `check_device` specifies  whether it checks the device assignment of the

      std::string group_annotation = absl::StrCat(
          "p", std::to_string(parallel_group_idx), ":", std::to_string(i));
      if (auto parallel_group_attr =
              parallel_execute_op->getAttrOfType<StringAttr>(
                  TF::kParallelExecAnnotation)) {
        // Extend the existing attribute so that nested parallel execution
        // structure is supported.

// Constraint that attribute is string with value either "SAME" or "VALID"
def IsSameOrValid : AttrConstraint<
    CPred<"$_self.cast<StringAttr>().getValue() == \"SAME\" || " #
          "$_self.cast<StringAttr>().getValue() == \"VALID\"">,
    "'SAME' or 'VALID' paddings">;

def TFL_GetMirrorPaddingType : NativeCodeCall<
    "mlir::TFL::MirrorPaddingTypeAttr::get($_builder.getContext(), " #

    return arg_shape;
  } else {
    return std::get<TensorShape>(arg.shape);
  }
}

Status MaybeRewriteLayoutWithShardedShape(
    mlir::StringAttr sharding,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns,
    xla::Shape* shape) {
  if (!sharding) return absl::OkStatus();

  xla::OpSharding op_sharding;