inline Type TransposeLastTwoDims(Type type) {
  auto shaped_type = type.dyn_cast<ShapedType>();
  if (!shaped_type.hasStaticShape() || shaped_type.getRank() < 2) {
    return nullptr;
  }
  int rank = shaped_type.getRank();
  if (rank < 2) {
    return nullptr;
  }
  SmallVector<int64_t> new_shape(shaped_type.getShape().begin(),
                                 shaped_type.getShape().end());

  auto input_type = mlir::cast<RankedTensorType>(input.getType());
  SmallVector<int64_t, 4> output_shape;
  int size_of_splits;
  if (input_type.getRank() < axis || axis < 0) return failure();
  for (int i = 0; i < input_type.getRank(); ++i) {
    int64_t dim = input_type.getDimSize(i);
    if (i == axis) {
      if (dim % splits != 0) {
        return failure();
      }

        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());
    if (output_type == nullptr || !output_type.hasStaticShape()) return false;

    int64_t cols = 1;
    for (int i = 0; i < output_type.getRank() - 1; ++i) {
      cols *= output_type.getDimSize(i);
    }
    const int64_t cost_per_col = 2 * weight_type.getNumElements();

    *count = cost_per_col * cols;

    auto bias = op->getOperand(2);

    // Input rhs must be a constant with rank 2.
    if (constant.getType().getRank() != 2) return failure();

    // Create a tfl.transpose op that performs ZX transpose on `input`.
    auto create_z_x_transpose_op = [&](Value input) -> Value {
      RankedTensorType input_type =
          mlir::cast<RankedTensorType>(input.getType());
      const int input_rank = input_type.getRank();

  // Index for the broadcasted matrix.
  llvm::SmallVector<int64_t, 16> current_index(result_type.getRank(), 0);
  // Computes the new operand shape using the original shape and the broadcast
  // dimensions to match result shape.
  llvm::SmallVector<int64_t, 16> operand_new_shape(result_type.getRank(), 1);
  for (int i = 0; i < dimensions.size(); ++i) {
    operand_new_shape[dimensions[i]] = operand.getType().getDimSize(i);

    CPred<"$0.getType().cast<ShapedType>().hasRank() && "
          "$0.getType().cast<ShapedType>().getRank() <= " # n>>;

// Checks if the value has rank 'n'.
class HasRank<int n> : Constraint<
    CPred<"$0.getType().cast<ShapedType>().hasRank() && "
          "$0.getType().cast<ShapedType>().getRank() == " # n>>;

class FloatValueEquals<string val> : Constraint<CPred<
  "FloatValueEquals($0, " # val # ")">>;

  if (!(tensor_type && tensor_type.hasRank())) {
    return op->emitError()
           << "cannot convert op with unranked or non-tensor input type "
           << tensor_type << ".";
  }

  int rank = tensor_type.getRank();
  if (rank <= partition_dim) {
    return op->emitError() << "cannot partition " << first_operand_type
                           << " (rank = " << rank << ") along dimension "

        op_builder->create<mlir::arith::ConstantOp>(
            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>(

  return lhs.hasRank() && rhs.hasRank() && lhs.getRank() == rhs.getRank();
}

// Creates a compatible RankedTensorType where mismatched dimensions are
// replaced with dynamic sizes.
RankedTensorType GetCompatibleRankedTensorType(RankedTensorType lhs,
                                               RankedTensorType rhs) {
  assert(lhs.getRank() == rhs.getRank());
  llvm::SmallVector<int64_t, 4> dims;

    if (passthrough_op->hasTrait<OpTrait::IsTerminator>()) return failure();

    auto get_num_elements = [](RankedTensorType tensor) {
      int num_elements = 1;
      for (int i = 0; i < tensor.getRank(); ++i) {
        // Assume dynamic dim size as the dim size one.
        if (!tensor.isDynamicDim(i)) {
          num_elements *= tensor.getDimSize(i);
        }
      }
      return num_elements;
    };