}
  return std::nullopt;
}

ShapedType GetNhwcReturnTypeFromNchw(Operation* old_op) {
  auto composite_result_shape =
      mlir::cast<ShapedType>(old_op->getResults().front().getType()).getShape();
  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW
  output_shape[0] = composite_result_shape[0];
  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];

Type ShuffleRankedTensorType(Type type, ArrayRef<int64_t> permutation) {
  if (auto ranked_type = mlir::dyn_cast<RankedTensorType>(type)) {
    ArrayRef<int64_t> shape = ranked_type.getShape();
    assert(permutation.size() == shape.size());

    SmallVector<int64_t, 4> new_shape(permutation.size());
    for (size_t i = 0; i < permutation.size(); ++i)
      new_shape[i] = shape[permutation[i]];

        type, builder.getFloatAttr(element_ty, unique_index));

  if (auto qtype = mlir::dyn_cast<QuantizedType>(element_ty)) {
    mlir::RankedTensorType new_type = tensorflow::GetTypeFromTFTensorShape(
        type.getShape(), qtype.getStorageType());
    return DenseElementsAttr::get(
        new_type, builder.getIntegerAttr(qtype.getStorageType(), unique_index));
  }
  llvm_unreachable("unhandled element type");
}

    DenseFPElementsAttr attr;
    if (!matchPattern(weight_op->getResult(0), m_Constant(&attr))) {
      return failure();
    }

    // Get new shape.
    llvm::ArrayRef<int64_t> cur_shape = attr.getType().getShape();
    int cur_rank = cur_shape.size();
    if (cur_rank != 4 || cur_shape[2] == 1) return failure();
    TensorType new_shape = RankedTensorType::get(
        {cur_shape[0], cur_shape[1], 1, cur_shape[2] * cur_shape[3]},

    Value bcast_op_result = (*bcast_op)->getResult(0);
    auto bcast_op_result_type =
        mlir::cast<RankedTensorType>(bcast_op_result.getType());
    const ArrayRef<int64_t> bcast_shape = bcast_op_result_type.getShape();
    const TensorType new_bcast_op_result_type = bcast_op_result_type.cloneWith(
        bcast_shape, accumulation_quantized_element_type);
    bcast_op_result.setType(new_bcast_op_result_type);
  }

                tensorflow::GetTypeFromTFTensorShape(
                    {splat_elements_attr.getType().getRank()},
                    op_builder->getI64Type()),
                splat_elements_attr.getType().getShape()));
    mlir::arith::ConstantOp fill_value =
        op_builder->create<mlir::arith::ConstantOp>(
            const_op->getLoc(),
            DenseElementsAttr::get(
                tensorflow::GetTypeFromTFTensorShape(

      if (v_int > 7 || v_int < -8) {
        return failure();
      }
    }

    Builder builder(op.getContext());
    auto shaped_type =
        mlir::RankedTensorType::get(const_type.getShape(), builder.getI4Type());
    auto newAttr = DenseElementsAttr::getFromRawBuffer(
        shaped_type, mlir::cast<DenseElementsAttr>(op.getValue()).getRawData());
    rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, newAttr);

    return PartialTensorShape();
  }

  if (auto tensor_type = mlir::dyn_cast<mlir::RankedTensorType>(type)) {
    TensorShapeProto tensor_shape_proto;
    ConvertToTensorShapeProto(tensor_type.getShape(), &tensor_shape_proto);
    return PartialTensorShape(tensor_shape_proto);
  }

  // If type is not a RankedTensor or UnrankedTensor, it must be a scalar.
  // Empty TensorShape indicates a scalar.

    }
  }
  // Verify axisStats (optional) attribute.
  if (getAxisStats()) {
    if (!getAxis()) return emitOpError("axis must be specified for axisStats");

    auto shape = tensorArg.getShape();
    auto argSliceSize =
        std::accumulate(std::next(shape.begin(), *getAxis()), shape.end(), 1,
                        std::multiplies<int64_t>());

    auto axisStatsType = getAxisStats()->getShapedType();

      auto inputType = mlir::dyn_cast<RankedTensorType>(arg.getType());
      // Only rank 1 tensor is supported for now.
      if (!inputType || inputType.getRank() != 1) continue;
      auto shape = llvm::to_vector<4>(inputType.getShape());
      llvm::SmallVector<int64_t, 4> bounds(shape.begin(), shape.end());
      // Mark the dim as dynamic dim.
      shape[0] = ShapedType::kDynamic;
      auto extensions =