if (!mlir::cast<ShapedType>(lhs.getType()).hasStaticShape() ||
        !mlir::cast<ShapedType>(rhs.getType()).hasStaticShape() ||
        !mlir::cast<ShapedType>(cond.getType()).hasStaticShape()) {
      return rewriteOpWithDynamicInput(op, rewriter);
    }

    auto lhs_shape = mlir::cast<ShapedType>(lhs.getType()).getShape();
    auto rhs_shape = mlir::cast<ShapedType>(rhs.getType()).getShape();

using ::mlir::stablehlo::DotGeneralOp;

bool HasStaticShape(Value value) {
  auto shaped_type = mlir::dyn_cast<ShapedType>(value.getType());
  if (!shaped_type) return false;

  return shaped_type.hasStaticShape();
}

bool HasStaticShapeAtDims(Value value, const ArrayRef<int> dims) {
  auto shaped_type = mlir::dyn_cast<ShapedType>(value.getType());
  if (!shaped_type || !shaped_type.hasRank()) return false;

  for (auto dim : dims) {

      if (operands.size() != 1 || updates.size() != 1) return failure();

      ShapedType operand_type = mlir::cast<ShapedType>(operands[0].getType());
      ShapedType indices_type = mlir::cast<ShapedType>(indices.getType());
      ShapedType updates_type = mlir::cast<ShapedType>(updates[0].getType());

      Value new_updates = updates[0];

      // Can only convert with static shaped scatter.

Value ConvertDot(PatternRewriter& rewriter, Value lhs, Value rhs,
                 mhlo::DotDimensionNumbersAttr dot_dimension_numbers,
                 ShapedType result_type, mlir::Location loc) {
  auto lhs_type = mlir::cast<ShapedType>(lhs.getType());
  auto rhs_type = mlir::cast<ShapedType>(rhs.getType());
  const int lhs_rank = lhs_type.getRank();
  const int rhs_rank = rhs_type.getRank();
  ImplicitLocOpBuilder builder(loc, rewriter);

bool IsLastDimEqualToNumElements(Type type1, Type type2) {
  return (mlir::cast<ShapedType>(type1).getRank() >= 1 &&
          mlir::cast<ShapedType>(type1).getDimSize(
              mlir::cast<ShapedType>(type1).getRank() - 1) ==
              mlir::cast<ShapedType>(type2).getNumElements());
}

bool CanFuseConvOrDepthwiseConvShapes(const ArrayRef<int64_t> filter_shape,

  if (auto bool_attr = mlir::dyn_cast_or_null<BoolAttr>(attr)) {
    return bool_attr.getValue();
  }
  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];

void BroadcastBatchDimensionsForBatchMatMul(OpBuilder &builder, Location loc,
                                            Value &input, Value &weight) {
  ShapedType input_type = mlir::cast<ShapedType>(input.getType());
  ShapedType weight_type = mlir::cast<ShapedType>(weight.getType());
  const int32_t input_rank = input_type.getRank();
  const int32_t weight_rank = weight_type.getRank();

    return builder.create<TF::ReshapeOp>(
        loc, RankedTensorType::get(shape, builder.getI32Type()), value,
        CreateConstValue<int64_t>(builder, loc, {rank}, shape));
  };

  ShapedType filter_shape = mlir::cast<ShapedType>(filter.getType());
  Value input_shape_value = builder.create<TF::ShapeOp>(
      loc, RankedTensorType::get({num_dims}, builder.getI32Type()), input);

      if (dim == axis) {
        if (ShapedType::isDynamic(operand_dim_size) ||
            ShapedType::isDynamic(result_dim_size)) {
          result_dim_sizes[axis] = ShapedType::kDynamic;
        } else {
          result_dim_sizes[axis] += operand_dim_size;
        }
        continue;
      }

      if (ShapedType::isDynamic(operand_dim_size)) continue;

      if (ShapedType::isDynamic(result_dim_size)) {

    Operation *op = srcop.getOperation();
    bool allTypesFp = true;
    bool allTypesQuantizedOrInt = true;
    for (auto operand : op->getOperands()) {
      ShapedType type = mlir::dyn_cast<ShapedType>(operand.getType());
      if (!type) continue;
      allTypesFp &= !mlir::isa<quant::QuantizedType>(type.getElementType());
      allTypesQuantizedOrInt &=