Jitting Stochastic Flax Module With Different Output Shapes

[moabarar]

Hi,

TL;DR: How to implement a stochastic Flax Module where the output shape is not deterministic?

I am trying to implement a stochastic Flax Module. For the sake of simplicity, I want to implement a module that performs the following:

• with probability 0.5, the module returns the inputs sliced along the last axis.
• otherwise, the module returns the inputs as is

I first implement something like this:

import flax.linen as nn
from jax import random

class MyStochasticModel(nn.Module):
  
  @nn.compact
  def __call__(self, inputs, rng = None):
    if rng is None:
      rng = self.make_rng('MyStochasticModel')
    condition = random.randint(rng, (1,), 0, 2)
    if condition[0] == 1:
      return inputs[...,0:inputs.shape[-1]//2]
    return inputs

To initialize the module and applying it I did the following:

rng = random.PRNGKey(0)
model = MyStochasticModel()
params = model.init({'MyStochasticModel':rng}, jnp.ones([10,10]))

To call the model:

model = MyStochasticModel()
stoch_model_rng = random.split(random.PRNGKey(0))[0]
model.apply(params, jnp.ones([10,10]), rngs={'MyStochasticModel' : stoch_model_rng}).shape

This works fine when calling the Module with different RNGs.

However, 'jitting' module didn't work:

(1) first approach was to directly jit the module as suggested in this issue

@jax.jit
def apply_fn( params, inputs, rngs):
  return MyModel().apply(params, inputs, rngs=rngs)

apply_fn(params, jnp.ones([10,10]), rngs={'MyStochasticModel' : stoch_model_rng})@jax.jit
def apply_fn( params, inputs, rngs):
  return MyModel().apply(params, inputs, rngs=rngs)

apply_fn(params, jnp.ones([10,10]), rngs={'MyStochasticModel' : stoch_model_rng})

But I get the following error:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 

(2) From the error I understand that I need to use jax.lax control-flow operations. But since the output shape is different, then I cannot simply use jax.lax.cond. See my attempt below:

class MyStochasticModel(nn.Module):
  @nn.compact
  def __call__(self, inputs,rng = None):
    if rng is None:
      rng = self.make_rng('MyStochasticModel')
    condition = random.randint(rng, (1,), 0, 2)
    return lax.cond(condition[0] == 1, lambda x: x[...,0:x.shape[-1]//2], lambda x: x, inputs)

But this doesn't work since the output shape for both branches should be the same. I receive the following error:

TypeError: true_fun and false_fun output must have identical types, got
ShapedArray(float32[10,5])
and
ShapedArray(float32[10,10]).

NOTE: In my case I want to use pmap - but I figured since the error is due to compilation issues so resolving this for jit would resolve my issue.

Thanks in advance for any help

[jheek]

This is something that fundamentally JAX is unable to do at the moment due to XLA being a static shape compiler. There are 2 options to work around this limitation:

1. Keep all shapes static. For example you can use padding/masking instead of dynamic shapes to "emulate" dynamic slicing.
2. Only JIT parts of the compute that have a static shape so you can for example jit a function like this:

def f(inputs, branch):
   if branch == 1:
     return x[...,0:x.shape[-1]//2]
   return inputs

f_jitted = jax.jit(f, static_argnums=(1,))  # compiles a seperate version of f for each value of branch
branch = random.randint(rng, (), 0, 2)
f(inputs, branch)