Can JAX handle irregular functions that produces tensors of different shapes?

[ifsheldon]

Hi everyone! I've been using PyTorch but I'm new to JAX.

I'd like to ask if JAX can handle irregular functions that produces tensors of variable shapes? This may seem to be weird at the first sight, but a real use case of mine is some simulation, in which the shapes of tensors have geometric meaning.

For the sake of simplicity, take a look at the below (insane) example.

import torch
def weird_algorithm(x):
    if x > 3:
        return x
    else:
        return torch.randn(3, 3) * x 

input_xs = torch.arange(5, requires_grad = True)
outputs = map(weird_algorithm, input_xs)
loss = torch.zeros(1)
for output in outputs:
    loss += output.sum()
loss.backward()

To make this example less insane, you can think of another example, which is training a fully-convolutional auto-encoder on images of different sizes without paddings.

This example is over-simplified, but you can see it's differentiable and a pure function, despite the algorithm generates irregular tensors of different shapes. In theory, it's totally fine to do the compute graph construction in parallel with Pytorch's Autograd, but I cannot really do the map in parallel with it because of GIL, so I wonder if JAX can do this and gives me correct gradients. Thanks a lot!

[jakevdp]

The answer is sort of. So long as the value that the if/else is branching on is a static value, it's fine. If the value is dynamic, then JAX would not support this within transforms such as JIT or vmap or autodiff. See this for more discussion of control flow in JAX: https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#control-flow

[mattjj]

Actually, autodiff (both forward and reverse mode, and arbitrary higher-order compositions) works great with control flow, dynamic shapes, and all that. It's jit and vmap which don't.

The original Autograd project, to which JAX is a direct successor, was the first Python library to get autoidiff of Python control flow "right", and it influenced the designs of the AD systems to follow (like Chainer, PyTorch, TF2, etc). With JAX we didn't want to give up any of that dynamism!

[jakevdp]

Confirmed, it works!

import jax
from jax import random

def weird_algorithm(x):
    if x > 3:
        return x
    else:
        return random.uniform(random.PRNGKey(0), shape=(3, 3)) * x

print(jax.jacobian(weird_algorithm)(1.0))
# [[0.35490513 0.60419905 0.4275843 ]
#  [0.23061597 0.6735498  0.43953657]
#  [0.25099766 0.27730572 0.7678207 ]]
print(jax.jacobian(weird_algorithm)(4.0))
# 1.0

[mattjj]

Here's a fairly direct translation of the original example:

import jax
import jax.numpy as jnp

import numpy as np  # used for np.random here

def weird_algorithm(x):
  if x > 3:
    return x
  else:
    return np.random.randn(3, 3) * x

def loss(input_xs):
  outputs = map(weird_algorithm, input_xs)
  # just `return sum(outputs)` would be better, but following the original here
  loss = jnp.zeros(())
  for output in outputs:
    loss += output.sum()
  return loss

input_xs = jnp.arange(5.)
g = jax.grad(loss)(input_xs)
print(g)

The main changes are:

1. Made a loss function to be differentiated. JAX autodiff APIs are about differentiating functions, just like mathematical differentiation is about functions. See the JAX Autodiff Cookbook.
2. Using numpy.random for a familiar side-effecting random number generator, which works fine with autodiff.

However, the random number generation here isn't really idiomatic JAX code. JAX has a functional PRNG design which works better with parallelization, advanced autodiff (e.g. rematerialization), JIT compilation, and vectorized batching (see this design note). It is less familiar though. (Another issue with using numpy.random as above is that it'd be generating random arrays on the CPU!)

Here's a more JAX-idiomatic way to express the same computation:

import jax
import jax.numpy as jnp

def weird_algorithm(key, x):
  if x > 3:
    return x
  else:
    return jax.random.normal(key, (3, 3)) * x

def loss(key, input_xs):
  keys = jax.random.split(key, 3)
  outputs = map(weird_algorithm, keys, input_xs)
  return sum(o.sum() for o in outputs)

key = jax.random.PRNGKey(0)
input_xs = jnp.arange(5.)
g = jax.grad(loss, argnums=1)(key, input_xs)
print(g)

Because of the if x > 3 part, we can't stage out this computation with jax.jit to be end-to-end optimized with XLA. But we can always apply jax.jit to smaller sub-functions. You can think of using JAX autodiff + jax.numpy as getting what other systems give you, with a fixed library of fast kernels to call into. But going beyond that you can use jax.jit to compile your own optimized kernels, so long as those compiled subroutines satisfy the restrictions that jax.jit requires. Luckily you can always compose together those optimized kernels with relatively unrestricted Python.

[ifsheldon]

Thanks for the detailed answers! It's very informative.

And I'd like to ask a few more questions........

1. Will map(weird_algorithm, keys, input_xs) be executed in parallel?
   I suppose not? It's unfortunate that my use case contains a lot of parallelizable pure mapping functions like the above example, which takes in a tensor of variable shapes and produces a tensor of variable shapes. And they are computationally intensive.

   One potential reason why I want to use JAX is that maybe jit can get around GIL (like torch.jit can) and parallelize this differentiable map(weird_algorithm, keys, input_xs) for me (which torch.jit cannot).

2. Can we jit the weird_algorithm with control primitives like lax.cond?
   As said by @mattjj ("It's jit and vmap which don't.") and mentioned in the doc-control flow, we cannot simply use if/else of Python, which is understandable. And then I saw lax.cond, which seems good enough.

   I think this is an obvious yes, but as I say in 1., my goal is not about JIT essentially but getting around GIL with JIT and parallelizing or vectorizing these expensive mapping functions.

3. Can we use instances of our custom classes, pass them to a jitted function and get gradients on some of their attributes?
   As I skimmed through the tutorials and the documentation, it seems JAX is supposed to be data/tensor-oriented, but in my physical simulation, there're so many tensors, so it would be more convenient to manage them with some meaningful classes.
   For example, an imaginary class will be Particle, like

   class Particle:
       mass: float # no need for gradient
       velocity: np.ndarray  # needs gradients
       property: np.ndarray # some properties, no need for gradients
   
   # can we jit this function?
   def calc_momentum(particle: Particle):  # could be method of Particle
       modified_momentum = particle.mass * particle.velocity + property
       return modified_momentum

   We could rewrite this function to make it only take tensors or scalars, but it's somewhat tedious.