auto transposed_type =
      RankedTensorType::get(transposed_shape, input_type.getElementType());
  DenseIntElementsAttr permutation = DenseIntElementsAttr::get(
      RankedTensorType::get(permutation_array.size(), rewriter.getI64Type()),
      permutation_array);
  return {permutation, transposed_type};
}

Value BuildIntConstOp(ImplicitLocOpBuilder& builder,
                      ConversionPatternRewriter& rewriter, int64_t const_value,

          Create1DConstValue<int64_t>(builder, loc, indices),
          Create1DConstValue<int64_t>(builder, loc, {index_map_size, 1})),
      /*value=*/
      builder.create<TF::CastOp>(
          loc,
          RankedTensorType::get(
              mlir::cast<ShapedType>(start_indices.getType()).getShape(),
              builder.getI64Type()),
          start_indices));

      alg = tensorflow::ConcreteRngAlgorithm::RNG_ALG_PHILOX;
    } else {
      return rewriter.notifyMatchFailure(op, "unsupported alg");
    }

    Type state_element_type = rewriter.getI64Type();
    RankedTensorType op_type = RankedTensorType::get(
        {tensorflow::RNG_MAX_COUNTER_SIZE + tensorflow::RNG_KEY_SIZE},
        state_element_type);
    if (op_type != rng_op.getType()) {

                                  rewriter.getI64Type()),
            start));
    auto size_attr = rewriter.create<TF::ConstOp>(
        value.getLoc(),
        DenseIntElementsAttr::get(
            RankedTensorType::get({static_cast<int64_t>(size.size())},
                                  rewriter.getI64Type()),
            size));

  Value input_batch_dims = builder.create<TF::SliceOp>(
      loc, RankedTensorType::get({num_input_batch_dim}, builder.getI64Type()),
      input_shape, zero, num_input_batch_dim_value);
  Value input_matmul_dims = builder.create<TF::SliceOp>(
      loc, RankedTensorType::get({num_matmul_dim}, builder.getI64Type()),
      input_shape, num_input_batch_dim_value, num_matmul_dim_value);

  Value weight_shape = builder.create<TF::ShapeOp>(

#include "tensorflow/lite/kernels/padding.h"

namespace mlir::quant {
namespace {

Value GetDimValue(OpBuilder &builder, Location loc, Value shape_value,
                  int32_t dim) {
  Type attribute_type = builder.getI64Type();
  return builder.create<TF::StridedSliceOp>(
      loc,
      RankedTensorType::get(
          {}, mlir::cast<ShapedType>(shape_value.getType()).getElementType()),
      /*input=*/shape_value,

            splat_elements_attr.getSplatValue<Attribute>()));
    auto broadcast_dims = DenseIntElementsAttr::get(
        RankedTensorType::get(/*shape=*/{0}, op_builder->getI64Type()),
        llvm::SmallVector<int64_t>{});
    mhlo::BroadcastInDimOp broadcast_in_dim_op =
        op_builder->create<mhlo::BroadcastInDimOp>(
            const_op->getLoc(), splat_elements_attr.getType(), scalar,

            const_op->getLoc(),
            DenseIntElementsAttr::get(
                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(),

    const SmallVector<int64_t, 4>& output_shape) {
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());
  auto indices_type = RankedTensorType::get(input_shape, builder->getI64Type());
  auto output_type = RankedTensorType::get(output_shape, builder->getF32Type());
  SmallVector<mlir::Type, 2> input_types{input_type, indices_type};
  SmallVector<mlir::Type, 1> output_types{output_type};

      loc,
      RankedTensorType::get(lhs_transposed_shape, lhs_type.getElementType()),
      lhs,
      DenseIntElementsAttr::get(
          RankedTensorType::get({lhs_rank}, rewriter.getI64Type()),
          lhs_permutation));

  // Transposes rhs shape to be in the order of {batch_dimensions, contracting
  // dimensions, out_dimensions}.
  llvm::SmallVector<int64_t, 4> rhs_permutation = Concat<int64_t>(