N-dim iteration

[AriMKatz]

Hi,

Is there a way to iterate over every element of a table without knowing the number of dimensions at function definition time? And generally, how would one write code that's dimension agnostic (if possible)?

(sorry if this is covered somewhere, I couldn't find it)

Thanks!

[duvenaud]

Great question. In some senses the answer is yes, in others it's no. Specifically, there's no way in the current type system to write a function that can take an arbitrarily-nested index set, like:

def reduce_over_arbitrary_dims (xs: m=>n=>o...=>y=>z=>Float) : Float = ...

However, there are three senses in which the answer is yes:

1. Nested index sets are isomorphic to products of index sets.

All for loops in Dex can already reduce over any index set that you can create, and index sets can be of arbitrary size and shape internally. For instance, here's the type of fsum:

def fsum {n} (xs:n=>Float) : Float = ...

As long as we can express your table indices as a single index set (here n), all reductions in Dex will already work.

How does this help us deal with nested tables? Concretely, nested tables of arbitrary shape m=>n=>o...=>y=>z=>a can be converted to and from flat tables of type (m & n & o ... & y & z)=>a for any particular shape. For example:

xs_table  : m=>n=>o=>Float = ...
xs_flat   : (m & n & o)=>Float = for (i, j, k). xs_table.i.j.k
xs_table' : m=>n=>o=>Float     = for i j k.     xs_flat.(i, j, k)

2. Tables themselves can be used as index sets

As of #876, tables themselves can be used as index sets, letting us build index sets whose dimension is determined at runtime, but still tracked by the compiler. For instance, the type (Fin D)=>letter represents the set of all D-letter strings. And the table for i:((Fin D)=>letter). i instantiates all those strings. For another example, ((Fin D)=>(Fin s))=>Float is a table of Floats with s^D elements. Note that the parentheses are important! (Fin a)=>(Fin b)=>c has a * b elements, but ((Fin a)=>(Fin b))=>c has b ^ a elements.

Constructing and summing over such an arbitrarily-dimensioned table could be done like so:

D = read_integer_from_file
arbitrary_dim_table : ((Fin D)=>(Fin s))=>Float = for i. func i
big_sum = sum arbitrary_dim_table

From the point of view of sum, its input is just a standard flat table of type n=>v. Here n happens to be (Fin D)=>(Fin s). But as always, no function or for loop that consumes an index set needs to know (or even can know) the internal structure of that index set.

This trick is used in the FFT demo to build vectors whose lengths are powers of 2 by definition, having type ((Fin log2_n)=>(Fin 2))=>Complex.

3. Record types might someday support iterating over fields.

As discussed in #258 and #308, it'd be useful in many contexts to support iteration over the fields of a record type. I think this would allow a stronger Yes to your question, and we could imagine writing e.g. a program that reads in a list of field names of arbitrary length and type, then reads a table with that many dimensions to populate it. The Fields type defined in the prelude can already support some related operations. However I think this poses some problems for the compiler and syntax. @apaszke has thought much more about these issues than me, and I get the impression he thinks it's possible.