state tx infer doesn't output predicted .X when using embed-key

[asalimih]

I have trained a state transition model which uses X_state embedding key and has a decoder to output the full genes. now when I want to use state tx infer and specify the embed-key, it doesn't put the pert_cell_counts_preds output of the decoder in the X of the final h5ad output of the infer.
Here is the command:

state tx infer \
--output /media/model/state_output \
--adata /media/test_template.h5ad \
--pert-col target_gene \
--checkpoint /media/model/state_output/checkpoints/step=64000.ckpt \
--embed-key X_state

[abhinadduri]

hi, thank you for making an issue.

to clarify, in your output, the .obsm['X_state'] key exists, but counts are never outputted?

[asalimih]

to clarify, in your output, the .obsm['X_state'] key exists, but counts are never outputted?

yes .obsm['X_state'] exists and contains the models latent ouput. .X (counts) also exists but it's the same as the input data. Not what the decoder outputs.

[tnnandi]

I had the same issue. It's probably because embed_key isn't passed as an input arg in (in main.py) run_tx_preprocess_infer(args.adata, args.output, args.control_condition, args.pert_col, args.seed). You need to add args.embed_key to this function call.

[asalimih]

I had the same issue. It's probably because embed_key isn't passed as an input arg in (in main.py) run_tx_preprocess_infer(args.adata, args.output, args.control_condition, args.pert_col, args.seed). You need to add args.embed_key to this function call.

Thanks but I think my issue is different than yours. here I'm using infer, not preprocess_infer

[bsimon11]

I had a similar issue where I wanted the embed_key to be something else (X_hvg) but to output into .X . I modified _cli/_tx/_infer.py to force the output to write to .X instead of --embed_key and it worked