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));
}

  // definition, potentially leading to bogus _replication_info attributes. So
  // we just scrub all tf.Constants of all extra attributes.
  // TODO(kramm): Remove this once tf.Const's folder is aware of extra
  // attributes.
  auto value_str_attr = StringAttr::get(&getContext(), "value");
  getOperation().walk([&](mlir::TF::ConstOp cst) {
    auto dict = cst->getAttrDictionary();

  return compile_mlir_result.status();
}

absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GraphToModule(
    bool unconditionally_use_set_output_shapes, const Graph& graph,
    llvm::ArrayRef<std::string> control_rets,
    const FunctionLibraryDefinition& flib_def, const GraphDebugInfo& debug_info,
    mlir::MLIRContext* context) {
  mlir::DialectRegistry registry;
  RegisterDialects(registry);

Christian Sigg <******@****.***> 1714640622 -0700