`TranslationProblem.prepare` method can't set `joint_attr={"attr":"X"}`

[selmanozleyen]

When I try to run this

tp = tp.prepare(batch_key="batch", src_attr=src_attr, tgt_attr=tgt_attr, joint_attr={"attr": "X"})

TemporalProblem will give an error because of this line https://github.com/theislab/moscot/blob/3151da7b324672fab6757b8a84bcd1d0bcb9f027/src/moscot/problems/cross_modality/_translation.py#L167 with the message KeyError: 'x_key'

[MUCDK]

Thanks @selmanozleyen , good catch. It happens because of this: https://github.com/theislab/moscot/blob/3151da7b324672fab6757b8a84bcd1d0bcb9f027/src/moscot/problems/_utils.py#L26

Now I am wondering whether this scenario is realistic, as the joint attribute in the case you describe would be the full gene space (no dim. reduction).

Still I think it's fair to implement it. The easiest way would be to do an if-else statement to get the shapes (i.e. where it fails atm). wdyt?

[selmanozleyen]

Now I am wondering whether this scenario is realistic, as the joint attribute in the case you describe would be the full gene space (no dim. reduction).

@MUCDK yeah I can do it, just wanted to check with the use case, if I should implement it or throw an error etc.