# Iterator IR improvements
## Overview
- Tuple of fields on different locations
```python
if some_condition:
field_a: Field[Vertex, ...] = ...
field_b: Field[Edge, ...] = ...
else:
field_a: Field[Vertex, ...] = ...
field_b: Field[Edge, ...] = ...
```
Lowering options:
Possible:
```
make_tuple(if_(...)(...)(...), if_(...)(...)(...), )
```
Not possible (as tuple of it not allowed):
```
if_(...)(make_tuple(it_a_tb, it_b_tb))(make_tuple(it_a_tb, it_b_tb))
```
- Multiple returns on different location types
_Very similar to tuple problem above._
```python
@field_operator
def field_op(inp_a: Field[[Vertex], ...], inp_b: Field[[Edge], ...]):
return inp_a, inp_b
```
Felix: Computation is representable on ITIR so not needed before
new version of IR.
- It maybe doesn't need to be represented on ITIR but in lower level performance IRs (e.g. gtfn IR?)
- Another option: moving both iterators to a fake location (`V_or_E`) with the bigger domain, and shifting both iterators to it with a fake "identity" connectivity and a mask/if
- Masked output
```python
@field_operator
def field_op(inp: Field[[Vertex], ...]):
return where(mask, inp+1, inp)
```
- Maybe the compiler can figure it out if there are not aliases for the variable. Check if it is possible to tell the compiler that there are not aliases or other similar hints.
- Reduce/partial shifts
Option 1: `reduce` takes a list of values and returns a value. `get_neighbor_values(...)` returns a (variable length) list.
Option 2: `reduce` takes an iterator. `neighbor_shift` returns an iterator of lists.
- Lowering of nested reduction (offsets in wrong order)
- Broadcast
```python
@field_operator
def field_op(...):
return broadcast(neighbor_sum(field, V2E), (V, V2E))
```
Proposed solutions:
```python
def stencil_like(*shifted_its) -> Value:
return make_tuple(deref(shifted_its)...)
def neighbor_sum(iterators):
return reduce(sum)(lift(stencil_like)(...))
```
Open question:
- What happens if deref invalid (optional values)
- Can we get size of tuple? Is this a list (just for reductions) instead of a tuple?
- How to distinguish between user asking for tuple argument in reductor.
- Direct access to sparse field neighbor.
- Does this solve the broadcast problem?
```python
@field_operator
def nb_sum_broadcast(field: Field[[C, C2V]]) -> Field[[C, C2V]]:
return broadcast(neighbor_sum(field, C2V), (C, C2V))
@field_operator
def some_op(field: Field[[C, C2V], ...]): -> Field[[C, C2V]]
return neighbor_sum(field, C2V)
@field_operator
def user_of_broadcasted(field):
return some_op(field_op(field))
```
```python
def nb_sum_broadcast(it): # it return a tuple of values on neighboring vertices
return reduce(plus)(it)
def some_op(nb_sum):
def foo_stencil(*its):
val = user_of_broadcasted(lift(neighbor_sum)(*its))
return deref(val)
```
- _
```python
@field_operator
def number_of_neighbors():
return neighbor_sum(broadcast(1, (C, C2V)), C2V)
# Possible solution ??
@field_operator
def number_of_neighbors():
auto_f = magic_ITIR_builtin_with_index_in_domain
return auto_f_as_ones * neighbor_sum(broadcast(1, (C, C2V)), C2V)
```
- Move static offset information to IR
- (learning type inference pass)
### ITIR Type Inference
Type inferrence on ITIR
[Core](https://github.com/GridTools/gt4py/blob/functional/src/functional/iterator/type_inference.py)