if (!op.get_XlaSharding().has_value()) return failure();

    NamedAttribute call_target_name = rewriter.getNamedAttr(
        "call_target_name", rewriter.getStringAttr("Sharding"));

    auto custom_call = rewriter.create<mhlo::CustomCallOp>(
        op.getLoc(), op.getType(), op.getInput(),
        ArrayRef<NamedAttribute>{call_target_name});
    custom_call->setAttr(kShardingAttr, op.get_XlaShardingAttr());

                                    concat_unary_operands.getResult(),
                                    op.getResult().getType(),
                                    ArrayRef<NamedAttribute>());
  Operation* new_unary_op = rewriter.create(new_unary_op_state);
  CopyDeviceAndUnderscoredAttributes(op, new_unary_op);

  rewriter.replaceOp(op, new_unary_op->getResults());

  return success();

  let hasRegionVerifier = 1;

  let builders = [
    OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$operands,
      "llvm::ArrayRef<::mlir::NamedAttribute>":$attributes,
      "unsigned":$numRegions),
    [{
      assert(numRegions == 2u && "mismatched number of regions");
      build($_builder, $_state, resultTypes, operands, attributes);
    }]>];

      ConversionPatternRewriter& rewriter) const final {
    if (op.getCallTargetName() != "ApproxTopK") {
      return mlir::failure();
    }
    auto is_supported_attr_name = [](NamedAttribute attr) {
      auto name = attr.getName();
      return name == "call_target_name" || name == "backend_config" ||
             name == "api_version" || name == "called_computations";
    };

  OperationState new_unary_op_state(loc, first_user->getName().getStringRef(),
                                    op.getOperand(), op.getOperand().getType(),
                                    ArrayRef<NamedAttribute>());
  Operation *new_unary_op = rewriter.create(new_unary_op_state);
  CopyDeviceAndUnderscoredAttributes(op, new_unary_op);

  // Unpack results after applying unary operation.
  auto unpack_unary_op =

        auto flex_builder = std::make_unique<flexbuffers::Builder>();
        size_t map_start = flex_builder->StartMap();
        auto attrs = vhlo_op->getAttrs();
        std::vector<mlir::NamedAttribute> extra_attributes;

        for (size_t i = 0; i < attrs.size(); ++i) {
          auto name = attrs[i].getName().str();
          auto attr = attrs[i].getValue();

        input.getType().cast<ShapedType>().getElementType());
    inferredReturnTypes.assign({result_type});
    return success();
  }

  // Default to non-time_major.
  std::optional<mlir::NamedAttribute> time_major_attr =
      attr.getNamed("time_major");
  bool time_majored =
      time_major_attr ? time_major_attr->getValue().cast<BoolAttr>().getValue()
                      : false;

  int64_t batch =