LogicalResult XlaVariadicSortOp::verify() {
  XlaVariadicSortOp op = *this;
  const auto &inputs_ty = op.getInputs().getType();
  int n_inputs = inputs_ty.size();
  auto input_ty_0 = inputs_ty[0].cast<ShapedType>();
  if (input_ty_0.hasStaticShape()) {
    for (int i = 1; i < n_inputs; ++i) {
      auto input_ty_i = inputs_ty[i].cast<ShapedType>();
      if (input_ty_i.hasStaticShape() &&

  DataFormatVecPermuteOp op = *this;
  auto input_ty = mlir::dyn_cast<RankedTensorType>(op.getX().getType());
  if (!input_ty) return success();

  int rank = input_ty.getRank();
  if (rank != 1 && rank != 2)
    return op.emitOpError("requires input of rank 1 or 2");

  if (rank == 1) {
    int64_t dim0 = input_ty.getDimSize(0);
    if (dim0 != ShapedType::kDynamic && dim0 != 4 && dim0 != 2)

    // as its an input requirement.
    if (!input_ty.hasRank() || input_ty.getRank() != 4) {
      return failure();
    }

    int64_t batch_cst = input_ty.getShape()[0];
    int64_t channels_cst = input_ty.getShape()[3];

    int64_t in_y_cst = input_ty.getShape()[1];
    int64_t in_x_cst = input_ty.getShape()[2];
    int64_t in_spatial_cst =

}

Type InferReductionOpType(Value input, Value reduction_indices,
                          BoolAttr keep_dims) {
  Type input_ty = input.getType();
  Type element_ty = getElementTypeOrSelf(input_ty);

  // Output type is unranked if input type is not ranked.
  auto ranked_ty = mlir::dyn_cast<RankedTensorType>(input_ty);
  if (!ranked_ty) return UnrankedTensorType::get(element_ty);
  int64_t rank = ranked_ty.getRank();

    }
    int axis = val.getSExtValue();

    auto input_ty = mlir::dyn_cast<RankedTensorType>(op.getInput().getType());
    if (!input_ty || !input_ty.hasStaticShape()) {
      return rewriter.notifyMatchFailure(
          op, "require the type of input to have static shapes");
    }
    ArrayRef<int64_t> input_shape = input_ty.getShape();
    int input_rank = input_ty.getRank();
    if (axis < 0) axis += input_rank;

  mlir::MLIRContext* context = tfr_module_.getContext();
  llvm::SmallVector<mlir::Type, 4> input_tys, output_tys;
  mlir::Builder builder(context);
  for (auto ty : input_dtys) {
    mlir::Type elt_ty;
    TF_RETURN_IF_ERROR(ConvertDataType(ty, builder, &elt_ty));
    mlir::TensorType mlir_ty = mlir::UnrankedTensorType::get(elt_ty);
    input_tys.push_back(mlir_ty);
  }
  for (auto ty : output_dtys) {
    mlir::Type elt_ty;

    // i.e., num_input/input_last_dim = num_result/k
    auto input_ty = mlir::dyn_cast_or_null<ShapedType>(value.getType());
    auto result_ty = mlir::dyn_cast<ShapedType>(slice_op.getType());
    if (!input_ty || !result_ty) return std::nullopt;
    if (!input_ty.hasStaticShape() || !result_ty.hasStaticShape())
      return std::nullopt;
    if (!input_ty.getRank() || !result_ty.getRank()) return std::nullopt;

  const RankedTensorType input_ty =
      input.getType().dyn_cast_or_null<RankedTensorType>();
  const RankedTensorType filter_ty =
      filter.getType().dyn_cast_or_null<RankedTensorType>();
  // If indeed both input type & filter type are ranked type and have ranks.
  // We will need to check their ranks are valid.
  if ((input_ty && input_ty.hasRank() && input_ty.getRank() != 4) ||

    private final ImmutableMap<E, Integer> inputSet;
    private final int mask;

    SubSet(ImmutableMap<E, Integer> inputSet, int mask) {
      this.inputSet = inputSet;
      this.mask = mask;
    }

    @Override
    public Iterator<E> iterator() {
      return new UnmodifiableIterator<E>() {
        final ImmutableList<E> elements = inputSet.keySet().asList();
        int remainingSetBits = mask;

tensor<*x!tf_type.resource>, %arg4: tensor<*x!tf_type.resource>) -> (tensor<?x?xf32>, tensor<*xf32>, tensor<?x16384xf32>, tensor<16x112x112x?xf32>, tensor<16x224x224x3xf32>, tensor<*xf32>) attributes {tf.entry_function = {control_outputs = "", inputs = "inputs_0,conv1_conv2d_readvariableop_resource,conv1_biasadd_readvariableop_resource,fc1000_matmul_readvariableop_resource,fc1000_biasadd_readvariableop_resource", outputs = "identity_RetVal,fc1000_matmul_readvariableop_RetVal,flatten_reshape_RetVal,relu_...