[rg111]

Deep Learning applications is one area.

Traditionally, OO code is written all the time.

But after I learned JAX/Flax, a light turned on inside my head and I now write functional Deep Learning code as much as I can. All my side projects and new code are purely functional in JAX/Flax. PyTorch has functional API known as functorch, and I have used it one project.

Where lots and lots of data in 3,4,5 dimensional tensors exist, and you need to run lots of transformation on them, and then you need to multiply a huge number of them thousand times in each second- functional code makes much more sense and gives immense sanity and peace of mind.

Those of you writing Deep Learning code, learn functional programming principles (immutable data, pure functions, leaving no side effect, etc.), and apply them to DL via functorch or JAX.

Your life will never be the same.

[debanjan16]

Do JAX and functorch have the same level of builtin functionalities (operations) as the original Pytorch library?

Where to learn about it other than the documentation?

[rg111]

JAX is very barebones and will require you to write much more code for the same task than you write in PyTorch.

Functorch is still new, and honestly, there is little to learn if you already know JAX. There are some talks from Meta, and then there is always the docs.