Error related to static arguments when using classes in JAX.

[srijaniiserprinceton]

Dear JAX experts,

I am a fairly new user of JAX but have been pleasantly surprised with the rewards of using it. I am trying to implement a class structure (yes, I have read the discussion in 1567 and the threads leading up to it and am trying out classes and namedtuples simultaneously).I run the following code

import jax
import jax.numpy as jnp
from collections import namedtuple

class qdptMode:
    def __init__(self, GVAR):
        """Initialized with parameters that do not change
        with the central multiplet CENMULT.
        """
        self.smax = GVAR.smax

    def get_mode_neighbors_params(self, CENMULT, NBR_DICT):
        """Updates the parameters corresponding to mode-neighbours.
        CENMULT is passed as a static variable since this update needs
        to happen only when CENMULT changes.
        """
        # THESE PARAMETERS OUGHT TO BE CONSTRUCTED ON THE FLY FROM CENMULT
        nl_neighbours = NBR_DICT.nl_neighbours + 1
        omega_neighbours = NBR_DICT.omega_neighbours + 1
        num_neighbours = jnp.size(NBR_DICT.omega_neighbours)

        NBR_DICT_FINAL = namedtuple('NBR_DICT_FINAL',
                                    'nl_neighbours \
                                    omega_neighbours \
                                    num_neighbours')
        
        NBR_DICT_FINAL = NBR_DICT_FINAL(nl_neighbours,
                                    omega_neighbours,
                                    num_neighbours)

        

        return NBR_DICT_FINAL
    

    def build_supermatrix_function(self):
        """Function that returns the function to calculate
        the superMatrix prior to solving eigenvalue problem. 
        This function is specific to the central multiplet. So,
        CENMULT needs to be a static argument.
        """
        def compute_supermatrix(CENMULT, NBR_DICT_INIT):
            """Function to assimilate all the neighbour info
            and return the function to compute the SuperMatrix'
            """
            # getting the neighbour dictionary from CENMULT
            NBR_DICT = self.get_mode_neighbors_params(CENMULT, NBR_DICT_INIT)

            # unpacking the neighbour dictionary
            # number of submatrices along each dimension of the square supermatrix
            dim_blocks = NBR_DICT.num_neighbours
           
            nl_neighbours = NBR_DICT.nl_neighbours
            
            # supermatix can be tiled with submatrices corresponding to                                                                                                      
            # (l, n) - (l', n') coupling. The dimensions of the submatrix                                                                                                    
            # is (2l+1, 2l'+1)                                                                                                                                               
            dimX_submat = 2*nl_neighbours[:, 1].reshape(1, dim_blocks) \
                          * jnp.ones((dim_blocks, 1), dtype='int32') + 1
            dimY_submat = 2*nl_neighbours[:, 1].reshape(dim_blocks, 1) \
                          * jnp.ones((1, dim_blocks), dtype='int32') + 1
            dim_super = jnp.sum(dimX_submat[0, :])
            # we use float32 for the current problem since for DR it is a real matrix
            supmat = jnp.zeros((dim_super, dim_super), dtype='float32')
            
            return supmat
            
        return compute_supermatrix

# initializing the necessary variables in gvar
fwindow = 150   # in muHz
# central multiplet
ell0 = 200
# max degree of perturbation (considering odd only)
smax = 5

# loading the neighbours. Here, we have hardcoded the following multiplets
# (0,198), (0,200), (0,202). 
nl_neighbours = jnp.array([[0,198], [0,200], [0,202]], dtype='int32')
omega_neighbours = jnp.array([67.65916460412984, 67.99455100807411, 68.32826640883721])

# this dictionary does not change with changing central multiplets
GVAR = namedtuple('GVAR', 'fwindow smax')

# this dictionary changes with central multiplet
# NEEDS TO BE A STATIC ARGUMENT
CENMULT = namedtuple('CENMULT', 'ell0')

# dictionary for neighbours
NEIGHBOUR_DICT = namedtuple('NEIGHBOUR_DICT', 'nl_neighbours omega_neighbours')

GVAR = GVAR(fwindow, smax)
CENMULT = CENMULT(ell0)
NEIGHBOUR_DICT = NEIGHBOUR_DICT(nl_neighbours, omega_neighbours)
# creating instance of the class qdptMode
qdpt_mode = qdptMode(GVAR)

# jitting function to compute supermatrix
_compute_supermatrix = jax.jit(qdpt_mode.build_supermatrix_function(), static_argnums=(0,))

__ = _compute_supermatrix(CENMULT, NEIGHBOUR_DICT).block_until_ready()

And come across the error

jax._src.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=0/1)> (https://jax.readthedocs.io/en/latest/errors.html#jax._src.errors.TracerArrayConversionError)

I am not sure why this should be happening. The error is on the line 82 "supmat = jnp.zeros((dim_super, dim_super), dtype='float32')". I declare CENMULT as a static argument since the dimension of "supmat" depends on CENMULT. And NEIGHBOUR_DICT will change only when CENMULT changes. So, I was under the impression that should work.

Any suggestions?

[jakevdp]

Thanks for the question!

When you run

dim_super = jnp.sum(dimX_submat[0, :])

within a JIT-compiled context, the result is a traced value regardless of whether the input is static. The sizes of matrices cannot depend on traced values, so when you run jnp.zeros((dim_super, dim_super)) you get the tracer conversion error.

I haven't looked closely at your full algorithm, but in order to use it within JAX transforms like JIT, you will have to write it in such a way that array sizes are static: one useful pattern is to use standard numpy ops when dealing with static values.

For more information on this, the How to think in JAX document is a useful read.

[srijaniiserprinceton]

AH! I see. I was actually confused for the longest time as to why the above example fails but why the following simple code works

import jax
import jax.numpy as jnp

def func(ell):
    m_arr = jnp.arange(-ell, ell+1, dtype='int32')
    return m_arr

# creating and compiling the jitted function
_func = jax.jit(func, static_argnums=(0,))
__ = _func(5).block_until_ready()

# running for different value of the static argument
for ell in range(1, 10):
    m_arr = _func(ell).block_until_ready()
    print(m_arr)

I tried out a bunch of things and realized as long as the scalars defining the dimensions of arrays that are created in jitted functions in JAX aren't traced values it works. In the above example, it isn't a traced value. But in my original post it got converted to a traced value after getting acted upon by jax.numpy operations.

Now your suggestions make a lot of sense. Just re-iterating so that I know I got this correctly: Any such scalar which determines the dimensions of arrays should not be traced values. So, throughout the code, to ensure they don't get converted to traced values, we need to keep applying numpy operations instead of jax.numpy operations. Is this correct?

My final question would then be: In such cases, is there a chance that we might be compromising the speedup that JAX gives us if there are numpy operations (in my mind are these speed breakers)? Would the code be slowed down due to such operations even when there is no re-compilation (meaning the cases where the static arguments are not changing)?

[jakevdp]

Any such scalar which determines the dimensions of arrays should not be traced values. So, throughout the code, to ensure they don't get converted to traced values, we need to keep applying numpy operations instead of jax.numpy operations. Is this correct?

Yes, that's correct – anything that affects static array properties (e.g. shapes and dtypes) must be static, i.e. not traced.

In such cases, is there a chance that we might be compromising the speedup that JAX gives us

All static operations are computed at trace-time, so you should see these as part of the JIT compilation cost. If you're applying the function multiple times to compatible inputs without changing the static arguments, you shouldn't see that cost at runtime.

[srijaniiserprinceton]

Thanks!