if (isa<TF::TPUCopyWithDynamicShapeOp>(op)) {
      if (failed(CheckOpIsValid(op))) return WalkResult::interrupt();
      if (CanMove(op)) {
        tpu_copy_with_dynamic_shape_ops.push_back(op);
      }
    }
    return WalkResult::advance();
  });
  if (walk_result.wasInterrupted()) return signalPassFailure();

  for (Operation* op : tpu_copy_with_dynamic_shape_ops) {

}

bool ShouldMoveOutOfRegion(
    Operation *op, Region *region,
    const llvm::DenseSet<ResourceHandle> &read_only_vars) {
  return ResourceOpCanBeHoisted(op, region, read_only_vars) ||
         (isMemoryEffectFree(op) && isSpeculatable(op));
}

bool OnlyHasReadEffect(Operation *op) {
  auto interface = llvm::dyn_cast<MemoryEffectOpInterface>(op);
  if (!interface) return false;

  for (auto& op : *block) {
    EvaluateCost(&op);
  }
}

void CostAnalysis::EvaluateCost(mlir::Operation* op) {
  if (auto cost_function =
          mlir::dyn_cast<tfrt::compiler::CostFunctionInterface>(op)) {
    cost_map_[op] = cost_function.cost();
    return;
  }

  if (!llvm::isa<mlir::TF::TensorFlowDialect>(op->getDialect())) {
    cost_map_[op] = max_arg_size_;

  builder->setInsertionPoint(op);
  auto operands = llvm::to_vector<4>(op->getOperands());
  operands.push_back(operand);
  auto new_op = builder->create<OpT>(op->getLoc(), op->getResultTypes(),
                                     operands, attr.getAttrs());
  op->replaceAllUsesWith(new_op.getOperation()->getResults());
  op->erase();
  return new_op;
}

// Returns success if the function has at most one op of the template type and

          op->getLoc(), op.getOperand().getType(), adaptor.getOperand());
      rewriter.replaceOpWithNewOp<mlir::stablehlo::BitcastConvertOp>(
          op, op.getResult().getType(), output);
    } else {
      // Both input/output are f32. Convert to no-op.
      rewriter.replaceOp(op, adaptor.getOperand());
    }
    return success();
  }
};
}  // namespace

#define GEN_PASS_DEF_CONVERTFUNCTOBFLOAT16PASS

        op.getLoc(), op.getType(), op.getOperands(),
        builder.getI64ArrayAttr(partition_dims), builder.getBoolAttr(false),
        op.get_XlaShardingAttr());
    op->replaceAllUsesWith(pi);
  } else {
    auto po = builder.create<TF::TPUPartitionedOutputV2Op>(
        op.getLoc(), op.getResultTypes(), op.getOperand(),
        builder.getI64ArrayAttr(partition_dims), op.get_XlaShardingAttr());

namespace mlir {
namespace TFL {
namespace tac {

// Returns true if 'op' is non const op. Returns false otherwise or if
// 'op' is null.
inline bool IsNonConstOp(Operation* op) {
  if (!op) return false;
  if (llvm::isa<arith::ConstantOp, mlir::func::ConstantOp>(op)) return false;
  if (op->hasTrait<OpTrait::ConstantLike>()) return false;
  if (llvm::isa<TFL::ConstOp, TFL::QConstOp>(op)) return false;
  return true;
}

// Checks if an op is inside a lifted function.
// If the given op pointer is a nullptr, returns false.
bool IsInLiftedFunc(Operation* op);

// Checks if the op is inside a StableHLO op with region.
// If the given op pointer is a nullptr, returns false.
bool IsInStableHloOpRegion(Operation* op);

// Checks if a given einsum op is supported for XlaDotV2 quantization.

        return true;
      };

  return RewriteOp(op, rewriter, get_broadcasted_shape);
}

template <typename Op>
LogicalResult ConvertResultsBroadcastableShapeOp::RewriteEqOp(
    Operation* op, PatternRewriter& rewriter) const {
  auto eq_op = llvm::dyn_cast_or_null<Op>(op);
  if (eq_op && eq_op.getIncompatibleShapeError())
    return RewriteOp(op, rewriter, OpTrait::util::getBroadcastedShape);
  return failure();

        op.getLoc(), /*resultType0=*/new_conv_output_tensor_type,
        /*lhs=*/input_transpose_op,
        /*rhs=*/filter_transpose_op,
        /*window_strides=*/op.getWindowStridesAttr(),
        /*padding=*/op.getPaddingAttr(),
        /*lhs_dilation=*/op.getLhsDilationAttr(),
        /*rhs_dilation=*/op.getRhsDilationAttr(),
        /*window_reversal=*/op.getWindowReversalAttr(),