RuntimeError: mat1 and mat2 shapes cannot be multiplied during sampling with finetuned_replogle.ckpt

[Bunnybeibei]

Hello! I'm trying to run the sampling script using the provided finetuned_replogle.ckpt.

It seems this checkpoint was trained on 12626 genes, but script defaults to a 2000-gene space, causing a shape mismatch error:
RuntimeError: mat1 and mat2 shapes cannot be multiplied (128x4000 and 25252x512)

I tried changing the config to match 12626 genes:

data.pad_length=12626
model.hidden_num=[12626,512]
model.input_dim=12626
data.embed_key=X

However, the script still fails because 2000 is strictly hardcoded in several assert statements across the sampling codebase (e.g., sampling_generation_helpers.py line 58 and sampling_generation.py line 94).

How can I correctly run sampling for this finetuned checkpoint? Are there plans to make the gene dimension dynamic instead of hardcoded?

Thanks!

[KatarinaYuan]

Hi,
Thanks for the question and sorry that I was too busy for the last few weeks.

For pretraining, to leverage different datasets, we build a unified gene vocabulary (size of 12626). That's why finetuned ckpt for Replogle is of shape 12626. This is similar for PBMC. (Note that Tahoe100M finetuned ckpt is an exception, because we replaced its input/output projection layer with 2000-dim one duering finetuning to avoid out-of-memory issue.)

The sampling logic for Replogle is that, we first sample 12626-gene expressions for all perturbed cells. And then, we only keep those 2000 HVG expressions for benchmarking. "sampling_generation_helpers.py line 58 and sampling_generation.py line 94" is loading the 2000 HVG pickle files to make sure the gene order for testing does not change. Therefore, these are expected coding logics.

Could you please paste your error logs in full details onto the github issue. Happy to help around.