bool prefer_tf2xla) {
    legalize_chlo_ = legalize_chlo;
    prefer_tf2xla_ = prefer_tf2xla;
    use_tf2xla_fallback_ = tf2xla_fallback_device_type.has_value();
    if (tf2xla_fallback_device_type.has_value()) {
      device_type_ = tf2xla_fallback_device_type.value().str();
    }
  }
  /// Performs the lowering to XLA dialect.
  void runOnOperation() override;
};

  auto type = result.getType();
  result.setType(converter.convertType(result.getType()));
  auto new_value = converter.materializeArgumentConversion(
      rewriter, result.getLoc(), type, {result});
  rewriter.replaceAllUsesExcept(result, new_value, new_value.getDefiningOp());
}

// Wrap operands in an an unrealized cast to create a cast to buffer any type

  auto reducer =
      builder->create<Op>(loc, block->getArgument(0), block->getArgument(1));
  builder->create<ReturnOp>(loc, reducer.getResult());
}

ConstantOp GetScalarConstOfType(Type ty, Location loc, int64_t raw_value,
                                OpBuilder* builder);

ConstantOp GetScalarNegZeroOfType(Type ty, Location loc, OpBuilder* builder);

// Converts an ArrayAttr to a 1D 64-bit dense elements attribute.

  auto lhs_map_or = EquationToMap(lhs);
  if (!lhs_map_or.has_value()) return std::nullopt;
  auto lhs_map = lhs_map_or.value();

  auto rhs_map_or = EquationToMap(rhs);
  if (!rhs_map_or.has_value()) return std::nullopt;
  auto rhs_map = rhs_map_or.value();

  auto out_map_or = EquationToMap(out);
  if (!out_map_or.has_value()) return std::nullopt;
  auto out_map = out_map_or.value();

      mlir::Value new_value = input.getOutput();

      if (mlir::isa<TF::TPUAnnotateTensorsWithDynamicShapeOp>(
              result.getDefiningOp())) {
        TF::TPUAnnotateTensorsWithDynamicShapeOp annotate_op =
            builder.create<TF::TPUAnnotateTensorsWithDynamicShapeOp>(
                op->getLoc(), result.getType(), new_value,
                result.getDefiningOp()->getAttrs());

  };

  auto op_has_device = [&](Operation& op, InferenceDeviceType& device) -> bool {
    std::optional<InferenceDeviceType> op_device =
        GetInferenceDeviceTypeForOp(&op);
    if (!op_device.has_value()) return false;
    device = op_device.value();
    return true;
  };

  auto op_device_is = [&](Operation& op, llvm::StringRef device) -> bool {
    InferenceDeviceType device_type;

    if (!IsOpWithQuantizableTrait(op) ||
        FindQuantizationUnitFromLoc(op->getLoc()).has_value()) {
      return failure();
    }

    std::optional<QuantizationUnit> quantization_unit =
        FindQuantizationUnit(op);
    if (!quantization_unit.has_value()) return failure();

    if (quantization_unit->func_name().empty()) {
      std::string func_name =

  return false;
}

// Check if op or its ancestor uses a key in `virtual_devices`.
bool AncestorUsesVirtualDevice(
    const std::optional<DictionaryAttr>& virtual_devices, Operation* op) {
  if (!virtual_devices.has_value()) return false;
  if (!op) return false;
  if (llvm::isa<tf_device::ReplicateOp>(op)) return false;
  if (DirectUseOfVirtualDevice(*virtual_devices, op)) return true;

    // Inversely transposes all the attributes to prepare for the new reduce op
    auto getInvertPermutedAttr =
        [&](std::optional<ArrayRef<int64_t>> vals) -> DenseI64ArrayAttr {
      return vals.has_value()
                 ? DenseI64ArrayAttr::get(
                       ctx, InvertPermutationToVector(*vals, transpose_perm))
                 : nullptr;
    };
    DenseI64ArrayAttr new_window_dimensions =

        if (input_sharding.has_value() && output_sharding.has_value() &&
            input_sharding.value() != output_sharding.value()) {
          return absl::InvalidArgumentError(absl::StrCat(
              "arg#", arg_idx, " is aliased to retval#", retval_idx,
              " but their sharding configurations don't match."));
        } else if (input_sharding.has_value() && !output_sharding.has_value()) {