[DaCe] Fixing Jax to SDFG Issues For a Comparison

# [DaCe] Fixing Jax to SDFG Issues For a Comparison ## Problem This document should be seen as an extension of [[DaCe] ITIR vs. Field View – The Best Representation](https://hackmd.io/@gridtools/BynVNnAu6), that focuses on the JAX to SDFG translator and how it should be fixed. Therefore, it is not a real shaped project. ##### Inability to Fuse Maps One of the issue we want to solve here is that some Maps, that are generated by the translator, are not put in the main Map. ![Resulting SDFG when using the `Jaxpr to SDFG` translator; especially the unfused Maps at the top.](https://hackmd.io/_uploads/r1zP8fFda.png) In the above SDFG a big part of the stencil was merged into a big Map. Because of that some large chunks of auxiliary memory have to be allocated. While this can be done in the beginning, it nevertheless leads to a larger memory footprint and more accesses. ##### Missing `scan` Operator The [`scan` primitive](https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.scan.html) is offered by Jax and there is no correspondence is not a `XLA` primitive. The translator currently does not implement it. ##### Missing `scatter` Primitive Translating the NumPy references to Jax will create a `scatter` operation, which the translator currently does not implement. ## Appetite - Inability to Fuse Maps: As this is not a blocker per sé, it is technically unneeded, however, I would like to invest 1 or 2 days to understand this issue a bit better. - Missing `scan` Operator: Will not be done. - Missing `scatter` Primitive: This is a must do part, however, it will come naturally as we integrate more stencils and should take 1 day. ## Solution ##### Inability to Fuse Maps There is no clear solution yet. It seems that the `CompositeFusion` transformation could do it, but for some reasons is not applied. Another possibility would be to modify the NumPy reference a bit should that other computational patterns appear. However, it is definitely a DaCe issue. It is important that this is not a real blocker. Therefore, we should not loose too much time with it.  ##### Missing `scan` Operator As it was mentioned in the [main document](https://hackmd.io/@gridtools/BynVNnAu6) scans are ignored for the time being. Thus we simply does not implement it. ##### Missing `scatter` Primitive According to my experience implementing the `scatter` primitive (Jax substitute every write access to a NumPy array with such an operation), will take a long time and should thus be avoided. However, most such accesses are used in the references to remove some computations that were performed by the reference but not by GT4Py, i.e. they have the form `A[:, -1] = 0`. Thus we will simply remove them, for verification these cleanups will be added again. ## Rabbit holes A lot. ## No-gos Solving other issues of the translator. ## Progress