// Returns a value if the `value` is defined by a ConstOp with a single
  // integer element in it and has an expected rank.
  auto get_const_int = [](Value value,
                          int expected_rank) -> std::optional<int64_t> {
    auto const_op = dyn_cast_or_null<ConstOp>(value.getDefiningOp());
    if (!const_op) return std::nullopt;

    auto value_attr = const_op.getValue().dyn_cast<DenseIntElementsAttr>();

    if (!op->hasAttr(kMiniBatchSplitsAttr) && !op->hasAttr(kMiniBatchCsrAttr))
      return WalkResult::advance();
    Operation* defining = op->getOperand(0).getDefiningOp();
    if (llvm::dyn_cast_or_null<TF::ConstOp>(defining)) {
      op->emitError("Couldn't find a program key to insert into this op.");
      return WalkResult::interrupt();
    }
    return WalkResult::advance();
  });

    return success();
  }

  LogicalResult processConstantOp(
      SourceOp op, Operation* const_op, int input_index,
      const operator_property::TensorProperty& tensor_property,
      PatternRewriter& rewriter) const {
    // Non-float tensors are neither weights nor require quantization.
    auto type = mlir::dyn_cast<ShapedType>(const_op->getResult(0).getType());
    if (!type || !mlir::isa<FloatType>(type.getElementType())) return success();

          squeeze_op, "result for current op has more than 1 use");
    }
    // Make sure that the axis in `expand_op` is constant.
    if (auto const_op =
            llvm::dyn_cast<TF::ConstOp>(expand_op.getDim().getDefiningOp())) {
      expand_axis = (*mlir::cast<DenseElementsAttr>(const_op.getValue())
                          .getValues<APInt>()
                          .begin())
                        .getSExtValue();

  ShapedType weight_type = mlir::cast<ShapedType>(weight.getType());
  const int32_t input_rank = input_type.getRank();
  const int32_t weight_rank = weight_type.getRank();
  const int32_t broadcasted_rank = std::max(input_rank, weight_rank);

  const int32_t num_matmul_dim = 2;
  const int32_t num_input_batch_dim = input_rank - num_matmul_dim;
  const int32_t num_weight_batch_dim = weight_rank - num_matmul_dim;

  for (const auto& arg : replicate.GetReplicatedBlockArguments()) {
    replicated_inputs.emplace_back();
    for (int64_t i = 0; i < num_replicas; ++i) {
      replicated_inputs.back().push_back(
          replicate.GetReplicaOperandForBlockArgument(arg, i));
    }
    new_replicated_inputs.emplace_back(replicated_inputs.back(), arg.getType());
  }
  for (const auto& arg : replicate.GetPackedBlockArguments()) {

}

bool IsTypeLegalizedWithMlir(const TypeID& type_id) {
  return MlirAlwaysOps().contains(type_id);
}

bool IsOpAllowedTf2xlaFallback(const TypeID& type_id) {
  return IsOpTypeAllowedTf2XlaFallback(type_id);
}

bool IsOpAllowedTf2xlaPreferred(const TypeID& type_id) {
  return IsOpTypeAllowedTf2XlaPreferred(type_id);
}

bool IsDynamicPadderOp(const TypeID& type_id) {

        return false;
    }
    if (!is_cluster_op) {
      if (!CheckNonClusterSuccessors(succ, op, metadata_map)) return false;
    }
  }
  return true;
}

bool TypeMustBeNonXLA(const Type& type) {
  const Type elem = getElementTypeOrSelf(type);
  return !mlir::isa<TF::ResourceType>(elem) &&
         !tensorflow::TypeValidForXLA(type);
}

IslandOp GetDummyConstant(OpBuilder builder, ShapedType const_type,
                          Location loc) {
  DenseIntElementsAttr val = DenseIntElementsAttr::get(const_type, 1);
  auto const_op = builder.create<TF::ConstOp>(loc, val);
  auto const_island = CreateIsland(const_op, {}, builder);
  return const_island;
}

// Rewrites the while op with extra chaining operands and results. Uses a

}

void TF_AttrBuilderSetTypeList(TF_AttrBuilder* builder, const char* attr_name,
                               const TF_DataType* values, int num_values) {
  auto iter = builder->attr_names.insert(attr_name).first;
  builder->Set(*iter, tensorflow::gtl::ArraySlice<const tensorflow::DataType>(
                          reinterpret_cast<const tensorflow::DataType*>(values),
                          num_values));
}