bug: `Identity` modality breaks using GPU backend

[federico-carrara]

Description

I recently tried to move to GPU backends (cupy) for simulating spectral data when simulating a small number of spectral bands (for large number of bins, I am always running into GPU OOMs).

I spotted a bug in the Identity modality at this line, as contrary to ground_truth, em_rates is never converted to the chosen backend before.

In my fork, I came up with this quick monketypatch:

em_rates = em_rates.sum(Axis.W)
em_rates = em_rates.copy(data=xp.asarray(em_rates.data))
em_image = em_rates * truth

@tlambert03 do you have a better alternative to propose?

[tlambert03]

using xp.asarray is the right way! that fix works for me. ... but do we need the copy? Why doesn't this work?

em_rates = xp.asarray(em_rates.sum(Axis.W).data)
em_image = em_rates * truth

I used to have a self-hosted cupy runner testing this, but that runner broke a while ago, so unfortunately it's hard to test these... i'll accept anything that you tell me works

[federico-carrara]

Hi Talley!

I missed your comment, sorry.
The reason why I do copy() is to keep em_rates a xr.DataArray. In this way, we can exploit xarray's smart broadcasting with auto-matching of dimensions for the subsequent tensor multiplication with truth.

If, instead, we do xp.asarray(em_rates.sum(Axis.W).data), em_rates becomes a numpy/cupy/etc. array, and we would have to add the logic to broadcast dimensions manually for the multiplication.