// If no such constant op exists, returns a nullptr.
Operation* GetBiasConstOp(Operation* op) {
  Operation* bias_const_op;
  if (Operation* broadcast_in_dim_op =
          FindOperandOfType<stablehlo::BroadcastInDimOp>(op);
      broadcast_in_dim_op != nullptr) {
    bias_const_op = broadcast_in_dim_op->getOperand(0).getDefiningOp();
  } else {
    bias_const_op = FindOperandOfType<stablehlo::ConstantOp>(op);
  }