Increase numerical precision for Jax (optional?) about moscot HOT 6 CLOSED

Ok, will add dtype to __init__ and ensure passed arrays/geometries have that dtype.

from moscot.

amazing, thanks!

from moscot.

I've tried 2 implementations:

1. user does

from jax.config import config
config.update("jax_enable_x64", True)

• we check and convert to f64 - problem is that all allocations (esp. ones not in our control in sinkhorn) are after the update in f64, which causes a mismatch between sizes.

2. add option x64: bool = False to fit - this worked better, but it's a stateful operation (I have to write to a config and then NOT revert it for it to work). If previously x64 was disallowed and was temporarily allowed for the computation, we again get some sizing erros (i.e. expected 1600 bytes, got 3200).
   Solution would be to convert back to x32 or to remain stateful and I dislike both (not to mention I don't know what the consequences of modifying the config like this are).

My take: if x64 is desired, use as jax does it in the beginning.

from moscot.

Mhm, that doesn't sound very promising, let's further discuss this today. How can I make sure I use float64 for my internal checks? I would like to run some checks.

from moscot.

Mhm, that doesn't sound very promising, let's further discuss this today. How can I make sure I use float64 for my internal checks? I would like to run some checks.

If you didn't use the above import, it should be always float32, otherwise float64.

from moscot.

Assuming this is done.

from moscot.