Question on slow JAX JIT compilation

[dhruvsreenivas]

Hi,

Hope you are doing well! I am working on tweaking a recent DeepMind algorithm (arXived in late June so official code is not out yet) for a new setting, and have implemented the loss function, but for some reason it takes quite a long time for JAX to JIT compile the code.

Attached is a code snippet that illustrates the loss function that I want to initialize. Basically, I am looping over a sliding window of a sequence of inputs, and pad with 0s appropriately (see comments in code for what I am trying to do). I'm not entirely sure what I am doing wrong in terms of making the code not really jittable.

def wm_loss_fn_window_size(self, wm_params: hk.Params, target_params: hk.Params, obs_seq: jnp.ndarray, action_seq: jnp.ndarray, window_size: int) -> Tuple[jnp.ndarray, jnp.ndarray]:

       seq_len = obs_seq.shape[0]
       obs_window = sliding_window(obs_seq, window_size).transpose(1, 0, 2, 3, 4, 5) # (window_size, num_seqs, B, 64, 64, 3)
       action_window = sliding_window(action_seq, window_size).transpose(1, 0, 2, 3) # (window_size, num_seqs, B, action_dim)
       num_seqs = action_window.shape[1]
       B = action_window.shape[2]

       # reshape to (window_size, B * num_seqs, *dims)
       obs_window = jnp.reshape(obs_window, (window_size, -1) + obs_window.shape[3:])
       action_window = jnp.reshape(action_window, (window_size, -1) + action_window.shape[3:])
       # print(f'window reshaped shapes: {obs_window.shape, action_window.shape}')

       pred_latents, _ = self.wm.apply(wm_params, obs_window, action_window)
       pred_latents = jnp.reshape(pred_latents, (-1,) + pred_latents.shape[2:]) # (T * B, embed_dim)

       _, target_latents = self.wm.apply(target_params, obs_window, action_window)
       target_latents = jnp.reshape(target_latents, (-1,) + target_latents.shape[2:]) # (T * B, embed_dim)

       # normalize latents
       pred_latents = l2_normalize(pred_latents, axis=-1)
       target_latents = l2_normalize(target_latents, axis=-1)

       mses = jnp.square(pred_latents - jax.lax.stop_gradient(target_latents)) # (window_size * B * num_seqs, embed_dim)
       loss = jnp.sum(mses, -1).mean()

       # take mean along the embedding dim
       mses = jnp.reshape(jnp.sum(mses, -1), (num_seqs, window_size, B))
       mses_padded = pad_sliding_windows(mses, seq_len)

       return loss, mses_padded

def wm_loss_fn(self, wm_params: hk.Params, target_params: hk.Params, obs_seq: jnp.ndarray, action_seq: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:

       seq_len = obs_seq.shape[0]
       init_val = (jnp.array(0.0), [])
       def body_fn(window_size: int, carry: Tuple[jnp.ndarray, List[jnp.ndarray]]):
               loss, mses_padded = self.wm_loss_fn_window_size(wm_params, target_params, obs_seq, action_seq, window_size)
               new_loss = carry[0] + loss
               new_carry = carry[1]
               new_carry.append(mses_padded)
               return new_loss, new_carry
       total_loss, seq_mses = jax.lax.fori_loop(1, seq_len + 1, body_fn, init_val)
       seq_mses = jnp.concatenate(seq_mses, axis=0)
       return total_loss / seq_len, jax.lax.stop_gradient(seq_mses)

The helper methods that are called in the above code snippet are in the snippet below. On my machine, the sliding_window function is jitted with no functools.partial on any argument and the pad_sliding_windows method is partially jitted with static seq_len parameter (although this may have to be changed due to unhashability of integers).

def sliding_window(arr: jnp.ndarray, window_size: int) -> Tuple[jnp.ndarray, jnp.ndarray]:

       idx = jnp.arange(arr.shape[0] - window_size + 1)[:, None] + jnp.arange(window_size)[None, :]
       return arr[idx]

def l2_normalize(x: jnp.ndarray, axis: Optional[int] = None, epsilon: float=1e-12) -> jnp.ndarray:

       square_sum = jnp.sum(jnp.square(x), axis=axis, keepdims=True) # from BYOL repo in deepmind-research
       x_inv_norm = jax.lax.rsqrt(jnp.maximum(square_sum, epsilon))
       return x * x_inv_norm

def pad_sliding_windows(arr: jnp.ndarray, seq_len: int) -> jnp.ndarray:

       arr = jnp.pad(arr, [(0, 0), (0, seq_len - arr.shape[1]), (0, 0)]) # pad with 0s first on right of dim 1
       rows, column_indices = jnp.ogrid[:arr.shape[0], :arr.shape[1]]
       roll_amts = jnp.arange(arr.shape[0])
       column_indices = column_indices - roll_amts[:, jnp.newaxis]
       return arr[rows, column_indices, :]

The specific error message is below. My jax/jaxlib version is 0.3.15 and my haiku version is 0.0.7.

Traceback (most recent call last):
  File "/Users/dhruvsreenivas/byol-offline/train.py", line 128, in main
    workspace.train_wm()
  File "/Users/dhruvsreenivas/byol-offline/train.py", line 88, in train_wm
    batch_metrics = self.wm_trainer.update(obs, action, self.global_step)
  File "/Users/dhruvsreenivas/byol-offline/byol_offline/models/wm.py", line 137, in update
    (loss, _), grads = loss_grad_fn(self.train_state.wm_params, self.train_state.target_params, obs, actions)
  File "/Users/dhruvsreenivas/byol-offline/byol_offline/models/wm.py", line 129, in wm_loss_fn
    total_loss, seq_mses = jax.lax.fori_loop(1, seq_len + 1, body_fn, init_val)
  File "/Users/dhruvsreenivas/byol-offline/byol_offline/models/wm.py", line 123, in body_fn
    loss, mses_padded = self.wm_loss_fn_window_size(wm_params, target_params, obs_seq, action_seq, window_size)
  File "/Users/dhruvsreenivas/byol-offline/byol_offline/models/wm.py", line 83, in wm_loss_fn_window_size
    obs_window = sliding_window(obs_seq, window_size).transpose(1, 0, 2, 3, 4, 5) # (window_size, num_seqs, B, 64, 64, 3)
  File "/Users/dhruvsreenivas/byol-offline/byol_offline/models/wm_utils.py", line 10, in sliding_window
    idx = jnp.arange(arr.shape[0] - window_size + 1)[:, None] + jnp.arange(window_size)[None, :]
  File "/Users/dhruvsreenivas/anaconda3/envs/byol-offline/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 2127, in arange
    start = require(start, msg("stop"))
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=3/1)>
It arose in jax.numpy.arange argument `stop`.
While tracing the function scanned_fun at /Users/dhruvsreenivas/anaconda3/envs/byol-offline/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py:1513 for scan, this concrete value was not available in Python because it depends on the value of the argument 'loop_carry'.

See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.ConcretizationTypeError

Any help would be greatly appreciated--I know that this might potentially be a dumb question with an easy answer but I'd love any help and explanation. I am also willing to explain what I was trying to do in more detail if need be. Thank you so much, I love the framework!

[jakevdp]

Would it be possible for you to paste the code in your question instead of pasting a screenshot? It's quite hard to help when it requires transcribing hundreds of lines of code from a PNG.

Thanks for updating the code snippets. From the error, it looks like the issue is you are tracing the window_size argument, when it needs to be static. Try something like this:

jitted_sliding_window = jax.jit(sliding_window, static_argnames=['window_size'])

then call this function on your inputs.

[jakevdp]

Yes, that's a possibility, but I suspect it will lead to long JIT compilation times because (as you mention) Python for loops are flattened and XLA will be compiling 50x as many instructions. Another possible approach for a situation like this is to use padding to avoid construction of dynamically-shaped arrays.

Note also that there is work in progress to make these kinds of situations easier (search "dynamic shapes" in JAX's github to see progress) but it's not fully implemented yet.

[dhruvsreenivas]

Thank you @jakevdp!

[dhruvsreenivas]

Hi @jakevdp, sorry for reopening this up, but I had a quick question on whether I could possible JAX JIT the following code:

def wm_loss_fn_window_size(self,
                               wm_params: hk.Params,
                               target_params: hk.Params,
                               obs_seq: jnp.ndarray,
                               action_seq: jnp.ndarray,
                               window_size: int) -> Tuple[jnp.ndarray, jnp.ndarray]:
        
        T, B = obs_seq.shape[:2]

        # define body function for given starting index 'idx' that goes into jax.lax.scan
        def body_fn(curr_state, idx):
            curr_loss, curr_loss_window = curr_state
            # get window starting from idx onward of obs, action
            obs_window = obs_seq[idx : idx + window_size] # (ws, B, *obs_dims)
            action_window = action_seq[idx : idx + window_size] # (ws, B, action_dim)
            
            pred_latents, _ = self.wm.apply(wm_params, obs_window, action_window) # (ws, B, embed_dim)
            pred_latents = jnp.reshape(pred_latents, (-1,) + pred_latents.shape[2:])

            _, target_latents = self.wm.apply(target_params, obs_window, action_window)
            target_latents = jnp.reshape(target_latents, (-1,) + target_latents.shape[2:])
            
            # normalize latents
            latents = l2_normalize(latents)
            embeddings = l2_normalize(embeddings)

            # take L2 loss
            mses = jnp.square(pred_latents - jax.lax.stop_gradient(target_latents)) # (ws * B, embed_dim)
            mses = jnp.reshape(mses, (-1, B) + mses.shape[1:])
            loss_vec = jnp.sum(mses, -1) # (ws, B)
            new_loss_window = curr_loss_window.at[idx : idx + window_size].add(loss_vec)

            return (curr_loss + loss_vec.mean(), new_loss_window), 0 # dummy stuff here

        init_val = (0.0, jnp.zeros((T, B)))
        (total_loss, loss_window), _ = jax.lax.scan(body_fn, init_val, jnp.arange(T - window_size))
        return total_loss, loss_window

@functools.partial(jax.jit, static_argnames=('self'))
def wm_loss_fn(self,
               wm_params: hk.Params,
               target_params: hk.Params,
               obs_seq: jnp.ndarray,
               action_seq: jnp.ndarray) -> Tuple[jnp.ndarray, jnp.ndarray]:
      
      T, B = obs_seq.shape[:2]
      def ws_body_fn(ws, curr_state):
          curr_loss, curr_loss_window = curr_state
          loss, loss_window = self.wm_loss_fn_window_size(wm_params, target_params, obs_seq, action_seq, ws)
          return curr_loss + loss, curr_loss_window + loss_window
      
      init_state = (0.0, jnp.zeros((T, B)))
      total_loss, total_loss_window = jax.lax.fori_loop(1, T + 1, ws_body_fn, init_val=init_state)
      return total_loss / T, total_loss_window / T # take avgs

This essentially takes the sliding window problem from before (where a dynamically-sized array was created) and instead does a loop over indices (I don't think you can do indexing with jax.lax.fori_loop, I saw on an older issue that you can maybe use jax.lax.scan which is what I did). Essentially, I wanted to create known, statically-sized arrays in the upper method, which we can then for loop over the indices via jax.lax.fori_loop.

I ran this, and came out with the following error:

File "/home/ds844/byol-offline/byol_offline/models/byol_model.py", line 190, in update
    (loss, _), grads = loss_grad_fn(self.train_state.wm_params, self.train_state.target_params, obs, actions)
  File "/home/ds844/byol-offline/byol_offline/models/byol_model.py", line 183, in wm_loss_fn
    total_loss, total_loss_window = jax.lax.fori_loop(1, T + 1, ws_body_fn, init_val=init_state)
  File "/home/ds844/byol-offline/byol_offline/models/byol_model.py", line 179, in ws_body_fn
    loss, loss_window = self.wm_loss_fn_window_size(wm_params, target_params, obs_seq, action_seq, T, ws)
  File "/home/ds844/byol-offline/byol_offline/models/byol_model.py", line 166, in wm_loss_fn_window_size
    (total_loss, loss_window), _ = jax.lax.scan(body_fn, init_val, jnp.arange(T - window_size))
  File "/home/ds844/.conda/envs/byol-offline/lib/python3.9/site-packages/jax/_src/numpy/lax_numpy.py", line 2136, in arange
    start = require(start, msg("stop"))
jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[], weak_type=True)>with<DynamicJaxprTrace(level=3/1)>
It arose in jax.numpy.arange argument `stop`.
The error occurred while tracing the function scanned_fun at /home/ds844/.conda/envs/byol-offline/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py:1599 for scan. This concrete value was not available in Python because it depends on the value of the argument 'loop_carry'.

My guess as to how to interpret this is that the sequence length parameter T is not static, but it should be known at compilation time right (given the definition above)? I also removed the jit over the total loss function wm_loss_fn and the same error persisted. Moreover, T in my case stays the same (static) throughout training, which makes things a bit weirder.

Let me know what you think. Thank you so much!

[jakevdp]

The issue is the same as above: you are constructing several arrays of size window_size; you would have to mark the window_size parameter as static to use this with JIT or other traced codepaths. Unfortunately, since ws is an argument to a fori_loop body function, it cannot be static, since its value changes with each loop.

[dhruvsreenivas]

I see, thanks! I did notice that you mentioned padding in the last response you gave, and I wanted to know how to best do it in this scenario--I don't know how best to pad the array properly (as I first create the dynamically sized window of the inputs themselves before I pass them through my neural networks).