---
# System prepended metadata

title: '[DaCe] Fixing Jax to SDFG Issues For a Comparison'
tags: [scratchpad]

---

# [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.
<!-- ![Hand Optimized SDFG.](https://hackmd.io/_uploads/HkdaF2zFT.png) -->


##### 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
<!-- Don't fill during shaping. This area is for collecting TODOs during building. As first task during building add a preliminary list of coarse-grained tasks for the project and refine them with finer-grained items when it makes sense as you work on them.

- [x] Task 1 ([PR#xxxx](https://github.com/GridTools/gt4py/pulls))
    - [x] Subtask A
    - [x] Subtask X
- [ ] Task 2
    - [x] Subtask H
    - [ ] Subtask J
- [ ] Discovered Task 3
    - [ ] Subtask L
    - [ ] Subtask S
- [ ] Task 4
-->
