if (!width) {
    return func.emitError() << "'width' attribute is not set or not an integer";
  }
  fbb.Int("width", width.getInt());

  auto string_separator =
      mlir::dyn_cast_or_null<StringAttr>(attrs.get("string_separator"));
  if (!string_separator) {
    return func.emitError()
           << "'string_separator' attribute is not set or not a string";
  }

  // "axis" operand could be a i64 tensor. Resolve it here.
  IntegerAttr axis_i32;
  if (failed(ConvertToI32Attr(axis_int, &axis_i32))) return failure();

  StringAttr fused_activation_function =
      StringAttr::get(rewriter.getContext(), "NONE");
  rewriter.replaceOpWithNewOp<ConcatenationOp>(
      op, output_type, values, axis_i32, fused_activation_function);
  return success();
}

  auto rendezvous_name = StringAttr::get(context, formatted_key);
  auto rendezvous_name_attr = NamedAttribute(
      StringAttr::get(context, xla::kXlaHostTransferRendezvousNameAttr),
      rendezvous_name);

  auto host_handler_name_value =
      StringAttr::get(context, host_handler_name.str());
  auto host_handler_name_attr = NamedAttribute(
      StringAttr::get(context, xla::kXlaHostTransferHandlerNameAttr),

            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,

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