# [GT4Py] NamedTuple-like Field Collections

(AKA pytrees)

- Shaped by:
- Appetite (FTEs, weeks):
- Developers: <!-- Filled in at the betting table unless someone is specifically required here -->

This task was originally shaped for [cycle 20](https://hackmd.io/4KDbjmE7SO6L_OjtLx2z0Q#). This project is stripped-down but with more details.

## Problem

GT4Py has basic support for collections via tuples. We identified several features that would profit from a general mechanism for collections.

See [cycle 20](https://hackmd.io/4KDbjmE7SO6L_OjtLx2z0Q#) for an overview.

In this project we aim to support

### NamedTuple-like collections

This is syntax sugar for long argument lists and allows to group related fields

Example: extraction

class Velocities:
    u: Field[DomT, DT]
    v: Field[DomT, DT]

def foo(velocity: Velocities) -> Field[[Dim], Dtype]:
    return sqrt(velocity.u**2 + velocity.v**2)

Example: creation

def foo(
    u: Field[[Dim], Dtype], v: Field[[Dim], Dtype]
) -> Velocities:
    return Velocities(u, v)

## Solution

With *pytree*s, JAX offers an extensible mechanism that provides this functionality. We would like to implement a similar mechanism for GT4Py.

A prototype is implemented in https://github.com/GridTools/gt4py/pull/2211 that sketches the ingredients of this project. Concrete decisions on how to expose registration and how to implement `flatten` (see later) are not yet taken.

### Access mechanisms

In this project we want to focus on attribute access for NamedTuple-like containers.

#### Expose type information at compile-time (parsing) and run-time
To support any user-defined container and not only a specific subset (e.g. NamedTuples, Dataclasses and Attr classes), the user needs to provide the following information:

- the GT4Py type of the collection: e.g. `NamedTupleType` that contains the `keys` (attributes) of the tuple elements on top of the plain tuple types.
- the GT4Py FunctionType for the constructor to be able to validate args/kwargs when constructing the object in the DSL.
- a flatten mechanism that converts the custom collection elements to a plain tuple which will be passed to the backend (this function is performance aka Python-overhead relevant), see [pytree protocol](https://jax.readthedocs.io/en/latest/pytrees.html) used by JAX or the the tree registering mechanism used in [autoray](https://github.com/jcmgray/autoray/blob/a0be3702f37bc64507ae501dd22717d224383ad2/autoray/autoray.py#L504).

### Lowering

We can resolve attribute access by `tuple_get` by mapping the attribute-name to the index as layed out in a `NamedTupleType`.

### Creation within the DSL (parsing)

The Custom container type need to be accessible including in type-checking. Utilities are already in place to get this information from ClosureVars.

## Rabbit holes
<!-- Details about the solution worth calling out to avoid problems -->

## No-gos

In this project we aim at basic support for NamedTuple-like attribute access. Implicit unrolling (expandable parameter style) and optimizations are not part of the project.

## 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