🎷JAX at NeurIPS 2020 - JAX on Cloud TPUs, JAX at DeepMind, JAX MD 🎷

[8bitmp3]

Hey core dev team @mattjj @skye @dougalm @hawkinsp @jekbradbury @gnecula @apaszke @jakevdp @zhangqiaorjc (did I miss anyone, @ Chris Leary?) 👋

Do you have any plans to present at NeurIPS or any other upcoming big ML conference any time soon? I noticed the Program Transformations for ML isn't in the NeurIPS workshop schedule this year https://github.com/program-transformations/program-transformations.github.io Maybe the workshop is on hold this year 🤔 cc @dynamicwebpaige

(I'm aware of a few recent JAX talks like https://www.youtube.com/watch?v=z-WSrQDXkuM by @jakevdp
and https://www.youtube.com/watch?v=lN3wWUdB_UM + https://www.youtube.com/watch?v=BzuEGdGHKjc by @mattjj and also this one https://www.youtube.com/watch?v=mbUwCPiqZBM by @froystig)

[mattjj]

Thanks for the question!

There may be a NeurIPS demo planned (cc @jekbradbury @skye).

I'm co-presenting a tutorial on deep implicit layers. That's not on JAX, but there is some JAX code involved!

[jekbradbury]

Wait, @8bitmp3, where are these chats from?

[8bitmp3]

@jekbradbury Rocketchat plugin that the organizers use for non-Zoom conversations - https://neurips.cc/virtual/2020/protected/tutorial_ad1df793247a0e650d0d7166341b8d97.html (login required). There was even a comment from Yoshua.

[mattjj]

Thanks, @8bitmp3 !

We spent a while polishing the captions, but somehow they didn't come through to the stream. The "on-demand" video has better (but less entertaining) captions.

[8bitmp3]

@jekbradbury

There was a comment on deep equilibrium models/DEQs' time to compute, in case you can't find it:

Q: If DEQs match the feedforward models' parameter counts, what's the compute time like?
A: For similar architectures/param numbers, "usually 2x-3x longer" because "the number of forward iterations we often find a suitably accurate fixed point is more than the "typical" depth of many networks"... but improving is "a direction of ongoing work"... "To clarify:... 1.5x at inference" (approx.).

Yoshua also said earlier that "a multi-layer DEQ may still help for the same computation resources. The same argument has been made for deep RNNs."

[jekbradbury]

We opened up a general JAX discussion channel on rocket.chat at #google-research-jax, in case anyone is watching this thread

[8bitmp3]

Reposting for visibility, since GitHub Discussions don't unroll some conversations automatically.

---

To sum up, here's the JAX-specific stuff at NeurIPS that you can find so far after the conference (UPDATED):

• December 8: JAX 💟 Cloud TPUs (11am-12pm PT/7pm-8pm GMT):

  • UPDATE: if you missed it - check the slides and 📺 https://twitter.com/jekbradbury/status/1337528357517291520 by @jekbradbury

• December 9: DeepMind 💙 JAX (Haiku, Optax, RLax, Jraph, etc) (11am-12pm PST/7pm-9pm GMT): (https://twitter.com/DynamicWebPaige/status/1336735477928345601 - by @dynamicwebpaige)

  • Blog post: https://deepmind.com/blog/article/using-jax-to-accelerate-our-research

• December 9/10: JAX MD: A Framework for Differentiable Physics (Orals & Spotlights Track 23: Graph/Meta Learning/Software, Poster Session 5) https://neurips.cc/virtual/2020/public/poster_83d3d4b6c9579515e1679aca8cbc8033.html / paper

also:

• December 7: Deep Implicit Layers: Neural ODEs, Equilibrium Models, and Differentiable Optimization with @mattjj (1.30pm-4pm PT, Tutorial Track 3C) https://neurips.cc/Conferences/2020/Schedule?showEvent=16653
  • JAX Colabs, YouTube 📺 recording, and more at http://implicit-layers-tutorial.org

(updated with Q&As from @jekbradbury, DeepMind talk public details from @dynamicwebpaige)