func_op->setAttr(
        kTfInputShapesAttr,
        ArrayAttr::get(func_op.getContext(), updated_input_shapes_attr));
  }
}

// Validates func ops. Returns `failure` if the function is invalid.
LogicalResult ValidateFuncOp(func::FuncOp func_op) {
  auto input_shapes_attr =
      func_op->getAttrOfType<ArrayAttr>(kTfInputShapesAttr);
  if (!input_shapes_attr) return success();

      !mlir::isa<IntegerType>(row_splits_type.getElementType())) {
    return func.emitError() << "Second input should be an integer tensor";
  }

  auto hash_seed =
      mlir::dyn_cast_or_null<ArrayAttr>(attr.getAttrs().get("hash_seed"));
  if (!hash_seed) {
    return func.emitError()
           << "'hash_seed' attribute is not set or not an array";
  }
  auto output_type = GetResultType(func, 0);

}

Status ConvertAttribute(const mlir::UnitAttr& attr, AttrValue* value) {
  value->clear_value();
  return absl::OkStatus();
}

Status ConvertAttribute(const mlir::ArrayAttr& attr, bool remove_ref_type,
                        AttrValue* value) {
  auto* list = value->mutable_list();
  for (mlir::Attribute a : attr.getValue()) {
    if (auto attr = mlir::dyn_cast<mlir::BoolAttr>(a)) {

        input_specs.push_back(TypeAttr::get(requantize.params));
      } else {
        input_specs.push_back(TypeAttr::get(state.params));
      }
    }
    op->setAttr("input_specs", ArrayAttr::get(context, input_specs));

    llvm::SmallVector<Attribute, 4> output_specs;
    auto original_output_specs = op.getOutputSpecs().getValue();
    for (int res = 0, e = op.getNumResults(); res != e; ++res) {

    auto replicate = cluster_func->getParentOfType<tf_device::ReplicateOp>();
    if (!replicate) return;
    auto mirrored_variable_indices_attr =
        replicate->getAttrOfType<ArrayAttr>(kMirroredVariableIndicesAttr);
    llvm::SmallDenseSet<int64_t, 8> mirrored_replicate_args;
    if (mirrored_variable_indices_attr) {
      for (const auto& mirrored_index : mirrored_variable_indices_attr) {

        *op_attrs) {
  assert(op_attrs);
  op_attrs->clear();

  mlir::Builder builder(context);
  for (auto iter : op_attr_array) {
    auto key_value = mlir::cast<mlir::ArrayAttr>(iter).getValue();
    llvm::StringRef key = mlir::cast<mlir::StringAttr>(key_value[0]).getValue();
    mlir::Attribute value = key_value[1];
    op_attrs->push_back({key, value});
  }
}

  }
  auto empty_array_attr = ArrayAttr::get(ctx, {});
  auto platforms = ArrayAttr::get(ctx, {StringAttr::get(ctx, kPlatformCpu)});

  auto call_op = builder.create<TF::XlaCallModuleOp>(
      location,
      /*output=*/output_types,
      /*args=*/args,
      /*version=*/kDefaultVersion, /*module=*/"",
      /*Sout=*/ArrayAttr::get(ctx, shape_attrs),
      /*dim_args_spec=*/empty_array_attr,

      if (!attr)
        return InvalidArgument(
            "Missing attribute '", output_arg.type_list_attr(),
            "' required for output '", output_arg.name(), "'");
      ArrayAttr array_attr = mlir::dyn_cast<ArrayAttr>(attr);
      if (!array_attr)
        return InvalidArgument("Attribute '", output_arg.type_list_attr(),
                               "' required for output '", output_arg.name(),

    return op->emitOpError()
           << "requires device ordinal from device " << logical_core_name
           << " to be present in 'tf.device.replicate' op";
  }
  llvm::StringRef tpu_device =
      tpu_replica.cast<ArrayAttr>()[replica_id].cast<StringAttr>().getValue();
  return tensorflow::GetDeviceOrdinalFromDeviceString(op->getLoc(), tpu_device,
                                                      &device_ordinal);
}

  }
}

mlir::LogicalResult GetDeviceAssignmentCoordinates(
    mlir::tf_device::ClusterOp cluster,
    llvm::SmallVector<int64_t, 8>& device_coordinates) {
  mlir::ArrayAttr device_assignment_attr =
      cluster->getAttrOfType<mlir::ArrayAttr>(
          tensorflow::kDeviceAssignmentAttr);
  if (!device_assignment_attr)
    return cluster.emitOpError(llvm::formatv("requires attribute '{0}'",