defn roadmap

[josevalim]

Constructs:

• Math operators
• Bit operators
• Logical operators
• Slices (access + put_in/update_in)
• Conditionals (if/cond)
• Tuples (with pattern matching)
• Loops (for? while?)
• Support random functions
• Default arguments

Passes:

• Autograd
• vmap
• pmap

[seanmor5]

Does it make sense to have pmap as its own pass/function/whatever, or does it make more sense to provide simple constructs for people to accomplish the same thing?

[josevalim]

To be honest, I don’t know yet, probably the latter. I just put it here so we don’t forget to track it but this is most likely a device/exla concern. I will move it. :)

[josevalim]

@seanmor5 so vmap may require changes to the underlying code to be compiled as it aims to add a new dimension to computations by making them batchable. So it is definitely a defn pass. It is hard to assess right now how big those changes are, since all of our operations are element wise so far. But with Jax, this code:

from jax import make_jaxpr

def f(x, y):
  a = jnp.dot(x, y)
  b = jnp.tanh(a)
  return b

xs = jnp.ones((8, 2, 3))
ys = jnp.ones((8, 3, 4))

print("f jaxpr")
print(make_jaxpr(f)(xs[0], ys[0]))

print("vmap(f) jaxpr")
print(make_jaxpr(vmap(f))(xs, ys))

prints:

f jaxpr
{ lambda  ; a b.
  let c = dot_general[ dimension_numbers=(((1,), (0,)), ((), ()))
                       precision=None ] a b
      d = tanh c
  in (d,) }
vmap(f) jaxpr
{ lambda  ; a b.
  let c = dot_general[ dimension_numbers=(((2,), (1,)), ((0,), (0,)))
                       precision=None ] a b
      d = tanh c
  in (d,) }

pmap, on the other hand, is about the devices. We want to move the data to separate devices when sending in and read them back from all multiple devices into a single binary. Or keep them on multiple references. So it is definitely a device based operation. I will create a separate issue for tracking the device roadmap.

[josevalim]

I have broken all remaining tasks to separate issues. The pmap discussion is tied to #127.