@tpoterba your method seems to be working well, though saving the results takes a really really long time. Is that to be expected?
batch_results = hl.linear_regression_rows(
[geno_subset_mt.phenos[i] for i in range(batch_size)],
geno_subset_mt.dosage,
covariates=[1.0],
block_size=128)
batch_results.write('test') # takes a really long time