Still facing NaN issues!

[chaithyagr]

Even after a lot of efforts, we still seem to face NaN issues and sadly this seems to be verry erratic and random. I dont have a repro test at all...
From what I see, I think that this issue is only present in Graph Mode. @jmontalt does that hint at any issues possible? Is there any separate code paths possible for eager mode / graph mode?

I will try to use this issue to track its progress on what I observe.

[jmontalt]

That is frustrating! Thank you for your efforts. Is this happening on the CPU or on the GPU?

Eager mode and graph mode follow the same code path. As far as I understand, the main difference between the two modes for our purposes is that in graph mode multiple ops may be executed concurrently. The errors might be related to this concurrent execution.

It might be worth trying to set tf.config.threading.set_inter_op_parallelism_threads to 1. If this makes the issue disappear, it would confirm the suspicion that it's related to concurrent op execution.

[chaithyagr]

Is this happening on the CPU or on the GPU?

Currently I am only working on GPU, so it is happening on GPU

It might be worth trying to set tf.config.threading.set_inter_op_parallelism_threads to 1. If this makes the issue disappear, it would confirm the suspicion that it's related to concurrent op execution.

Thats a good thing to test. I will try this. However, that still build the question as to how will we fix such an issue in case it is this. I atleast hope this will be faster to run as compared to eager mode..

[chaithyagr]

In all the runs I launched yesterday all ran fine for 20h these are:

1. Repeat an earlier run which resulted in NaN
2. Run in eager mode
3. Run with tf.config.threading.set_inter_op_parallelism_threads as 1

This is getting really frustating as I really just dont have a setup to repro to even be ...
All t he above runs was made with all the seeds set to 0. I am pretty much lost at this point.. But I will keep the issue open to track the presence of such an issue.. Is that fine @jmontalt ?

[jmontalt]

Hi @chaithyagr, thank you so much for running these tests. It's very useful to know this. Just to be clear, you have never been able to reproduce the issue if either running on eager mode or if inter-op parallelism threads is 1. Is that correct?

Yes, let's keep this issue open. Would you be willing to share some code that has failed before, even if it does not reproduce the issue consistently? It would give me a starting point to run some tests with code that has at least been known to fail.

Interestingly, I don't currently observe this problem, and at this point I'm not sure whether this could be due to differences in our usage or something else (perhaps it's less likely on my hardware?).

[chaithyagr]

Just to be clear, you have never been able to reproduce the issue if either running on eager mode or if inter-op parallelism threads is 1. Is that correct?

Correct, and I get your point where you expect this to fix our issue and frankly I dont get a huge speedup with inter-op parallelism set as something other than 1. (I am working in 3D, so I guess we are pretty much using all resources already!)
However, as the repeat run also passed, I am confused :P

Yes, let's keep this issue open. Would you be willing to share some code that has failed before, even if it does not reproduce the issue consistently? It would give me a starting point to run some tests with code that has at least been known to fail.

I think the same code from #20 (comment) should be good enough minimum repro. Although one problem I see with this version is that in TFMRI we now threhold the density (https://github.com/mrphys/tensorflow-mri/blob/cfd8930ee5281e7f6dceb17c4a5acaf625fd3243/tensorflow_mri/python/ops/traj_ops.py#L938) and this non_nan reciprocal could be bad:

out = tf.math.reciprocal_no_nan(tfmri.sampling.estimate_density(traj, img_size, method, max_iter))

Interestingly, I don't currently observe this problem, and at this point I'm not sure whether this could be due to differences in our usage or something else (perhaps it's less likely on my hardware?).

I dont think it could be hardware as I run it on multiple machines and different GPUs. It could be possible that something is wrong with my codes specifically. I now plan to save the input when we get NaNs (although this may not be of much use as an extact repeat of run, which results in same output, did not result in NaN, signyfying that it is a more random issue)

[chaithyagr]

A minor question, @jmontalt I was wondering what happens if the input points are not in [-pi, pi], could this possibly cause an errors / NaNs? Do we have any guard checks?

[jmontalt]

It should support up to the extended range [-3 * pi, 3 * pi], with points beyond [-pi, pi] simply being folded back to this range.

This extended range should give some flexibility and robustness against points that are close to [-pi, pi], but note that I have never tested this. It would be useful to add such a test.

For points beyond this extended range, behaviour is undefined.

We do not currently assert that the points lie within the correct range. This would be rather inefficient, especially on the GPU.

This behaviour is the same as in the original FINUFFT.

[chaithyagr]

Well, while i agree the wrapping on a mathematical point of view, I think this isnt the valid in MRI. But i guess tenorflow-nufft is generic package so I guess its fine.

[jmontalt]

It does remain valid in MRI. Note that pi is, by definition, the maximum spatial frequency (in radians per pixel) that you can have in an image of a certain size (which is one of the inputs to nufft). In order to usefully sample higher frequencies than that, you would need a finer grid. But if you have a finer grid, then pi is simply redefined to be the maximum spatial frequency for that grid.

[chaithyagr]

Yes, I get that. What i mean is that in the case of MRI if the inputs are given as is (not re-normalizing to pi), then:

with points beyond [-pi, pi] simply being folded back to this range.

This would be wrong.. But then again, as i told, I dont think thats the job of tensorflow-nufft as it is a more generic library..

[chaithyagr]

For points beyond this extended range, behaviour is undefined.

If this is valid. I think this could be the issue. I observed on adding a lot of randomness to my k-space trajectory, the NaN issue came up suddenly. While I still cant pin point why, but my repro has always been in controlled setting where there is no points outside the range. Lets try to see how the code behaves in this scenario. I will try to create such an odd repro for now

[chaithyagr]

Ok, I think that was the issue!
If I take the repro from above and basically add np.pi to the k-space locations, while NUFFT works fine, the gradients with respect to the trajectories are NaN!

Not sure if this is realistic though. But proper guard checks are added on my side, I will report back if the issue still remains.

[jmontalt]

Great, thanks for keeping me in the loop.

So forward nufft works as expected (e.g., with wrapping) with points in [-3 * pi, 3 * pi], but the gradient only works in [-pi, pi]? Or is it that you had points beyond [-3 * pi, 3 * pi]?

I'll make a note to add an option to assert that the points are in the right range, which would be useful for debugging purposes (I don't want to enable this by default because of likely performance hit).