Using custom data structure in JIT compiled code

[leviyevalex]

Hello!

An algorithm I want to implement requires efficient removal, selection, and random choice. A pure python data structure that accomplishes this task was discussed here for example.

Since JAX requires static shaped arrays, I've attempted to develop a class which accomplishes the same task using only JAX arrays. This is my first time using a custom data structure in JAX, and would appreciate any feedback on whether this will fail in any spectacular way.

import numpy as np
import jax.numpy as jnp
from functools import partial
import jax
from jax import tree_util

class ParticleSystem:
    """ 
    O(1) lookup, removal, and choice.
    """
    def __init__(self, N):
        self.N = N 
        self.array = jnp.zeros(shape=(N,2), dtype='int32')
        self.array = self.array.at[:,0].set(jnp.arange(N))
        self.array = self.array.at[:,1].set(jnp.ones(N, dtype='int32'))
        self.n_dead = 0
        self.pointer = jnp.arange(N)

    def kill(self, i):
        """ 
        Kill particle i. This involves the following:
        (1) Setting boolean value held in `array` to 0
        (2) Shifting dead particle entry to the end of `array`
        (3) Updating `pointer` to keep track of where entries are in `array`
        """
        self.n_dead += 1
        j = self.pointer[i]

        # Update "lookup" table
        self.array = self.array.at[j,1].set(0)

        # Swap entries in array
        b = self.array[self.N - self.n_dead]                                     
        self.array = self.array.at[self.N - self.n_dead].set(self.array[j])      
        self.array = self.array.at[j].set(b)                                     

        # Swap entries in pointer
        k = self.pointer[b[0]]                                                   
        self.pointer = self.pointer.at[b[0]].set(self.pointer[i])               
        self.pointer = self.pointer.at[i].set(k)                                 

    def choose_alive(self, key):
        """ 
        Randomly choose a particle that is alive. 
        This is efficient because particles are organized according to alive/dead in `array`.
        """
        randint = jax.random.randint(key, (1,), 0, self.N - self.n_dead)[0]
        return self.array[randint, 0]
    
    def is_alive(self, i):
        """ 
        Returns binary value, (0,1) (dead, alive)
        """
        return self.array[self.pointer[i], 1]

I am currently running some tests to see if this code behaves as expected. Using the discussion here, I have

def flatten_ParticleSystem(obj):
    children = [obj.array, obj.n_dead, obj.pointer]
    aux_data = (obj.N)
    return children, aux_data

def unflatten_ParticleSystem(aux_data, children):
    obj = ParticleSystem(aux_data)
    obj.array = children[0]
    obj.n_dead = children[1]
    obj.pointer = children[2]
    return obj

jax.tree_util.register_pytree_node(ParticleSystem, flatten_ParticleSystem, unflatten_ParticleSystem)

def body_fun(i, val): 
    val.kill(i)
    return val

@partial(jax.jit, static_argnums=(0,))
def test_func(nParticles): 
    a = ParticleSystem(nParticles)
    return jax.lax.fori_loop(3, nParticles, body_fun, a)

res = test_func(6)
print(res.array)
print(res.pointer)

which appears to yield the correct result

[[0 1]
 [1 1]
 [2 1]
 [5 0]
 [4 0]
 [3 0]]
[0 1 2 5 4 3]

Assuming the logic of the class is correct, may I use this class normally in my code moving forward? Are there any JAX specific quirks I'm not accounting for? For my use case, $N \sim \mathcal{O}(1000)$.

Suggestions are greatly appreciated!

[jakevdp]

Hi! This looks good, the only concerning thing I see is the use of in-place mutation in your class. For example, res.kill(0) will result in res being modified, but jax.jit(res.kill)(0) will not, because JAX requires pure functions.

As long as you don't write code that expects mutation side-effects to persist across transformation boundaries, it should be fine. For example, res.kill is impure and therefore problematic, but your body_fun as defined is non-problematic despite calling val.kill, because the mutations are reflected in the return value of the function.

[leviyevalex]

Could you elaborate on how body_fun is not problematic?

My actual method uses logic very similar to this, and doing things this way seems like a great way to keep the code clean and organized!

[jakevdp]

The reason body_fun is fine is because it does not rely on side-effects: it returns the updated object as an output rather than trying to modify the input in-place.

[leviyevalex]

Perfect! Thank you so much!