Allow user to specify `max_batch_size`? about tensorflow-nufft

[jmontalt]

Hi @chaithyagr, I think this is a good idea and I'd be happy to support the feature.

I would suggest that such a parameter goes into an Options structure which would then be passed to nufft (similar to FINUFFT itself). For example:

import tensorflow_nufft as tfft

opts = tfft.Options(max_batch_size=4)

result = tfft.nufft(source, points, options=opts)

There are a few reasons why I think this approach is a good idea:

• It makes it clear to users that these are advanced options that provide fine-grained control over the performance and memory usage, but otherwise do not have an effect on the result.
• It keeps the nufft signature and doc page simple for most users, and those with higher customization needs can delve deeper into the Options structure.
• It recognises that there are many low-level parameters that we might want to expose in the future (as indeed there are internally already), and it makes sense to keep them together.

A good model might be tf.data.Options. This is basically a Python structure which is serialized as a protocol buffer before passing to the OptionsDataset op.

Note that I'm not necessarily suggesting we should add any other internal parameters now, as each of them will involve some work adapting code and testing. But the Options structure provides a straightforward way to do so in the future.

Let me know what you think and whether you are happy to implement this!

[chaithyagr]

This is a great idea but for now, I seem to have implemented the basic version which directly feeds from nufft interp and spread functions. I do agree that this makes codebase and also the module a bit clunky right now. But i had to do it for my personal project. I can start this PR for now as we work on the Options Module as you mentioned, what do you think?

[jmontalt]

Thank you for your contribution. However, I am not in favour of making "temporary" changes to the public API. I'll keep the PR open as it's a good starting point, but I don't think it is ready to be merged yet. I'll keep this in mind and will do my best to help with the implementation when I find some time. In the meantime, I would suggest that you keep this version for your own work while we set up the final solution.

[chaithyagr]

Ya, I was also not very happy with my own solution. But I dont have time yet to work on this. It is a nice-to-have feature though. For now, I have my own version.

[jmontalt]

This has been addressed in #30 and the requested feature will be available in the next release.

[chaithyagr]

Wonderful, thank you for your efforts @jmontalt !

[jmontalt]

No problem! It should be easy to add more options in the future. Anything else you would like to add, please do not hesitate to ask (or create a PR directly).