Introduce while loops

[josevalim]

I would like to add a higher-level looping construct to defn. In Elixir, we can use for+:reduce but I am afraid it will be too verbose and foreign for new users. Let's imagine we want to translate this python code:

    def _smooth(x):
      out = np.empty_like(x)
      for i in range(1, x.shape[0] - 1):
          for j in range(1, x.shape[1] - 1):
              out[i, j] = x[i + -1, j + -1] + x[i + -1, j + 0] + x[i + -1, j + 1] +
                          x[i +  0, j + -1] + x[i +  0, j + 0] + x[i +  0, j + 1] +
                          x[i +  1, j + -1] + x[i +  1, j + 0] + x[i +  1, j + 1]) // 9
      return out

With for+:reduce, we would write it as:

    def smooth(x) do
      for i <- 1..elem(x.shape, 0)-1, j <- 1..elem(x.shape, 1)-1, reduce: x do
        x ->
          put_in x[i, j], x[i + -1, j + -1] + x[i + -1, j + 0] + x[i + -1, j + 1] +
                           x[i +  0, j + -1] + x[i +  0, j + 0] + x[i +  0, j + 1] +
                           x[i +  1, j + -1] + x[i +  1, j + 0] + x[i +  1, j + 1]) / 9
      end
    end

I propose we introduce a loop construct, inspired by futhark that looks like this:

loop tuple_or_var [= expr], [pattern <- expr]+ do

end

Rewriting the above to this loop construct, we have:

    def smooth(x) do
      loop x,
           i <- 1..elem(x.shape, 0)-1,
           j <- 1..elem(x.shape, 1)-1 do
        put_in x[i, j], x[i + -1, j + -1] + x[i + -1, j + 0] + x[i + -1, j + 1] +
                          x[i +  0, j + -1] + x[i +  0, j + 0] + x[i +  0, j + 1] +
                          x[i +  1, j + -1] + x[i +  1, j + 0] + x[i +  1, j + 1]) / 9
      end
    end

There is one downside with this approach: the only form of loops we have in XLA are while loops which are sequential. Other languages, such as taichi, can optimize them to run in parallel. For this reason, we may want to introduce higher level constructs for manipulating tensors. In particular, I believe we should introduce functions such as map, map_with_index, reduce and reduce_with_index. I have some thoughts in how we can implement said functions so they also work with batching out of the box, but I am waiting for some feedback on this issue before moving forward.

[seanmor5]

I am a fan of futhark so I think adding their style of looping is definitely an interesting idea. In terms of adding map, map_with_index, reduce and reduce_with_index, I am confident in being able to do that with EXLA.

Understanding that the for+:reduce concept might be foreign/unfamiliar to new users, I am curious if maybe it's a good idea to include both options for users who want their code to be closer to what a normal Elixir program might support?

[josevalim]

@seanmor5 that's a good question. I went back to Furhark code and I realized the support two looping constructs: for and while:

loop x while x >= 1
loop x for y in expr

If we want to support both loops, then it means this is not quite the same as for+:reduce. I think we can support while loops like this:

loop x, x >= 0 do
  x - 1
end

So I would rather proceed with loop and then if people think this is a good idea, we can even consider having it as part of Elixir?

[josevalim]

After thinking more about this, I suggest for us to start with while only and make it clear it is a sequential construct. We can explore something like for in the future.