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

Activity

1. KatarinaYuan commented on Apr 7, 2026

   

   KatarinaYuan

   on Apr 7, 2026

   Collaborator

   More actions

   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.
2. kiristern commented on May 20, 2026

   

   kiristern

   on May 20, 2026

   More actions

   Hi @KatarinaYuan,
   Thanks for this work!

   I am trying to simply reproduce the same figures as in your paper, using your provided Replogle dataset and checkpoints. I am coming up with the same dims mismatch issue as above when trying to run rawdata_diffusion_sampling.py on replogle finetune checkpoints.

   I made the following .sh modifications to try to run sampling on replogle finetune checkpoint:

   PRETRAIN_CKPT_PATH="${ROOT_PATH}PerturbDiff_ckpt/finetuned_replogle.ckpt"
   
   ./scripts/run1.sh
   
   REPLOGLE_SAMPLING_EXTRA="
   model_checkpoint_path='$PRETRAIN_CKPT_PATH'
   data=replogle_finetune
   data.selected_gene_file='${ROOT_PATH}PerturbDiff_data/perturb_data/selected_genes/merged_pbmc_tahoe_rep_cellxgene_genes_mapped.pkl'
   sampling.num_sampled_batches=null
   cov_encoding.celltype_encoding=llm
   model.replace_1w2gene_layer=true
   "
   
   python ./src/apps/run/rawdata_diffusion_sampling.py \
   $COMMON_SAMPLING \
   $REPLOGLE_SAMPLING_EXTRA \
   $COMMON_MODEL_12626 \
   $COMMON_REPLOGLE_TRAIN \
   $NO_WANDB
   

   where scripts/run1.sh is a direct copy-paste from your main/README.md.

   Note, I removed the data.sample_replogle_only=true because this appeared to change the path on line 54 of src/apps/sampling/sampling_generation.py to PerturbDiff_data/perturb_data/selected_genes/replogle_real_selected_genes.pkl (but it appears the finetune data should use the PerturbDiff_data/perturb_data/selected_genes/merged_pbmc_tahoe_rep_cellxgene_genes_mapped.pkl, if I understand correctly?).

   However, then the top 2000 HVG error is raised, which I think is due to missing huggingface files (line 32 of configs/cov_encoding/trixie_onehot.yaml):

   gene_embedding_path: ["${path.tmp_dir}/gene_names/pbmc_highly_variavle_gene_emb_dict_emb_dict.pkl",
                         "${path.tmp_dir}/gene_names/replogle_gene_emb_dict_perturbation_emb_dict.pkl",
                         "${path.tmp_dir}/gene_names/replogle_highly_variavle_gene_emb_dict_emb_dict.pkl",
                         "${path.tmp_dir}/gene_names/tahoe100m_highly_variavle_gene_emb_dict_emb_dict.pkl"]
   

   (Note assumed typos: variavle -> variable ?).

   I was only able to find and download replogle_gene_emb_dict_perturbation_emb_dict.pkl in perturb_data/gene_names. Will the missing top HVG.pkl files be uploaded in the future?

   Meanwhile, I tried to reproduce the replogle (from scratch) instead, using:

   PRETRAIN_CKPT_PATH="${ROOT_PATH}PerturbDiff_ckpt/from_scratch_replogle.ckpt"
   
   ./scripts/run1.sh
   
   REPLOGLE_SAMPLING_EXTRA="
   model_checkpoint_path='$PRETRAIN_CKPT_PATH'
   data=replogle_finetune
   data.sample_replogle_only=true
   data.selected_gene_file='${ROOT_PATH}PerturbDiff_data/perturb_data/selected_genes/replogle_real_selected_genes.pkl'
   cov_encoding.celltype_encoding=llm
   sampling.num_sampled_batches=null
   "
   
   python ./src/apps/run/rawdata_diffusion_sampling.py \
   $COMMON_SAMPLING \
   $REPLOGLE_SAMPLING_EXTRA \
   $COMMON_MODEL_2000 \
   $COMMON_REPLOGLE_TRAIN \
   $NO_WANDB
   

   This script appears to work, however after running cell-eval run ... on the output diffusion_predict.h5ad and diffusion_true.h5ad, I don't get the same values as in your paper Table 4. Additionally, when plotting a simple scanpy.pl.umap on the generated vs. real outputs, there is no overlap between points.... I must be doing something wrong.

   Would it be possible to please share your workflow for reproducing your paper's figures ? Thanks again for your time !
3. KatarinaYuan commented on Aug 23, 2026

   

   KatarinaYuan

   on Aug 23, 2026

   Collaborator

   More actions

   Hi @kiristern, thanks for the detailed report.

   I looked into this more carefully. The released finetuned_replogle.ckpt is a 12626-gene checkpoint, so sampling must also be run in the 12626-gene input space:

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

   However, please keep:

     data.sample_replogle_only=true           
   

   This flag is important. It does not only switch the selected-gene path; it also makes the datamodule sample only Replogle and makes the sampling output use the Replogle 2000 HVG set for benchmarking. The intended logic is:

   1. sample in the 12626 unified gene vocabulary;
   2. then restrict the generated and real expressions to the 2000 Replogle HVGs for evaluation.

   So the assert len(_genes) == 2000 in the sampling code is expected for the benchmark output stage.

   The problematic part in your command is removing data.sample_replogle_only=true while setting:

     data.selected_gene_file=.../merged_pbmc_tahoe_rep_cellxgene_genes_mapped.pkl
   

   That merged file contains 12626 genes and is used for the model input vocabulary, not for the final Replogle HVG benchmark output. For Replogle sampling, the final selected genes should be replogle_real_selected_genes.pkl.

   Also, the files with highly_variavle in the filename are unfortunately named with that spelling in the current data release/code, so the spelling itself is not the runtime issue.

   A correct Replogle finetuned sampling command should look like:

   python ./src/apps/run/rawdata_diffusion_sampling.py \
       $COMMON_SAMPLING \
       model_checkpoint_path=${ROOT_PATH}PerturbDiff_ckpt/finetuned_replogle.ckpt \
       data=tahoe100m_pbmc_replogle_pretrain_cellxgene \
       data.sample_replogle_only=true \
       cov_encoding.celltype_encoding=llm \
       cov_encoding.replogle_gene_encoding=genept \
       cov_encoding.replace_pert_dict=true \
       sampling.num_sampled_batches=null \
       data.pad_length=12626 \
       model.input_dim=12626 \
       'model.hidden_num=[12626,512,2048,500]' \
       data.embed_key=X \
       optimization.micro_batch_size=128 \
       data.use_cell_set=32 \
       $NO_WANDB
   

   I also found that the current release code/documentation does not make this workflow clear enough, and there were some sampling-side compatibility issues around loading the released checkpoint and overriding local embedding paths. I will update the README / sampling script so the Replogle finetuned checkpoint can be reproduced directly.

   With the corrected workflow, the generated output should have shape (n_cells, 2000) after HVG restriction, and the Replogle finetuned metrics should be close to the paper values.