Marking a non-jittable inner function as non-jittable.

[ozencgungor]

Hi there,

I've been wondering if there is a way in jax to mark (for the lack of a better word) an inner non-jittable function to avoid it raising errors when outer jax.jit or jax.pmap calls try to compile it. The full code would be too long to copy here but to give an idea:

import jax
import flax.linen as nn
from typing import Any

Array = Any
def some_non_jit_func(x, a, b):
    y = a*b
    return y[y >= x]

class inner_layer(nn.Module):
    param: Array
    @nn.compact
    def __call__(self, inputs):
        x = param*x

class wrapper_layer(nn.Module):
    a: int
    b: int
    x: int

    @nn.nowrap:
    def get_layer(self, param):
        return inner_layer(param)
    
    @nn.compact:
    def __call__(self, inputs):
        param = f(self.x, self.a, self.b)
        layer = self.get_layer(param)
        x = layer(param)
        return x

after which the wrapper_layer is used inside another nn.Module and is to be trained in a jax.pmaped training loop very similar to the imagenet example.

The helper function some_non_jit_func is probably non-jittable and in this case only gets called a couple times during flax.linen model initialization and it's not computationally too expensive but it is buried deep and it would take me significant rewriting to expose it in order to be able to mark some of its arguments as static. I figured in order to be able to do some testing before I try to make it jittable would be to find some way to force jax to not compile the function but I cannot seem to find a way to do that.

Any help would be greatly appreciated.

[YouJiacheng]

You can use host_callback.
Alternatively:

import jax
import flax.linen as nn
from typing import Any

Array = Any
def some_non_jit_func(x, a, b):
    y = a*b
    return y[y >= x]

def make_model(x, a, b):
  z = some_non_jit_func(x, a, b) # do anything
  class inner_layer(nn.Module):
      param: Array
      @nn.compact
      def __call__(self, inputs):
          return param * inputs
  
  class wrapper_layer(nn.Module):  
      @nn.nowrap:
      def get_layer(self, param):
          return inner_layer(param)
      
      @nn.compact:
      def __call__(self, inputs):
          layer = self.get_layer(z)
          x = layer(inputs)
          return x
  return wrapper_layer

[mattjj]

Thanks for the question!

You might be able to use jax.ensure_compile_time_eval, though to be sure it would help if you could provide a runnable toy example.

[ozencgungor]

Thanks for the replies! After some thinking, I think I might be able to come up with a jittable version of the function in question which leads into another question. The easiest way for me to do that would be to use fancy indexing on scipy.sparse.lil_array objects but Im expecting that to not go well in a jax.lax.cond. I think I'd rather figure out how to set slices of certain rows and columns of a BCOO array to zero in a jax.lax control flow statement.

I'll try and post a shortened workable example today.

[ozencgungor]

I'm marking this as answered as I figured the best approach is to figure out a way to efficiently set slices of rows of a BCOO array to zero. Thanks again for all the quick replies.

[ozencgungor]

As I suspected, it turns out it is indeed possible to construct a jittable version of my non-jittable function. For those who might be interested, here's what I was trying to do. I could not find a way to get rid of the jax.lax.fori_loopbut I'm happy with it for now.

import jax.numpy as jnp
from jax.experimental import sparse as jaxsparse

def block_padded_L_v2(L, nside, nside_sup, indices):
    """
    converts the original graph laplacian into a block diagonal laplacian in the following format
    
           |L_0  0   0   0...   0 |
           | 0  L_1  0   0...   0 |
           | 0   0  L_2  0...   0 |
        L= | 0                  0 |   where N is the number of patches(superpixels.)
           | 0   .  ...         0 |
           | 0   .              0 |
           | 0   0   0   0...  L_N|
           
    and L_n is the padded graph laplacian for patch n. 
    --------------------------------------------------
    L: the original graph laplacian in jax BCOO format
    nside: nside of the full map
    nside_sup: nside of the superpixel map, in the above expression N=12*nside_sup**2
    indices: indices of full map where masked indices are set to -1
    --------------------------------------------------
    returns: block diagonal L as in above, L for the superpixels    
    """
    def block_diag(L, block_size):
        """
        Turns M into block diagonal format where each block is of shape (block_size, block_size)
        by putting zeros in places that do not belong to a block
        ----------
        M: jax BCOO array
        block_size: size of blocks
        """
        data = L.data
        idx = L.indices
        row = idx[:,0]
        col = idx[:,1]
        block_num = row//block_size
        blocks_delim_low = block_num*block_size 
        blocks_delim_high = block_num*block_size + block_size
        data = (blocks_delim_low<=(col))*((col)<blocks_delim_high)*data
        new_L = jaxsparse.BCOO((data, idx), shape=L.shape)
        return new_L

    def index_mapping_fun(a, b):
        """
        a: COO indices (jaxified)
        b: unpadded indices (jaxified)
        """
        return jax.lax.fori_loop(0, a.shape[0], lambda i, x: x.at[i].set(b[x[i]]), a)

    def pad_L(L, nside, indices):
        """
        Pads a graph laplacian with rows and columns of zeros corresponding to masked pixels.
        --------
        L: sparse BCOO array
        nside: nside of the maps
        indices: unmasked indices
        --------
        returns BCOO array of shape (12*(nside**2),12*(nside**2))
        """
        data = L.data
        row = L.indices[:,0]
        col = L.indices[:,1]
        new_row = index_mapping_fun(row, indices)
        new_col = index_mapping_fun(col, indices)
        pad_L = jaxsparse.BCOO((L.data, jnp.column_stack((new_row, new_col))), shape=(12*(nside**2),12*(nside**2)))
        return pad_L    
    
    num_patches = 12*(nside_sup**2)
    block_size = 12*(nside**2)//num_patches
    padded_L = pad_L(L, nside=nside, indices=indices)
    block_diag_L = block_diag(padded_L, block_size)
    return block_diag_L

[VinaySingh561]

Hi @ozencgungor , could you please help me here?
I am using this function:

@jax.jit
def energy_fn(params,position):
    x,y,pos,force,edge_indexxxx= train_dataset[0]
    padded_edge_index = pad_edge_index(pos,500)
    data = (x,y,position,force,tuple(padded_edge_index))
    return new_e3nn_model.apply(params,random_key,data)+monomer_energy

and i want to make sure that third line of the code ""padded_edge_index = pad_edge_index(pos,500)"" does not get inside the jit compiled code as the function 'pad__edge_index' is using dynamic array.