Compilation when indices of vector depends on values stored in another vector (CSR matrix)

[dbalabka]

In the case of a sparse CSR matrix, indices of one vector are stored in another. For example, we have a sparse matrix of user-item interactions in CSR format, and we try to find if a person had an interaction with some particular item:

import jax.experimental.sparse as jsp
import jax.numpy as jnp

interactions = jsp.BCSR.fromdense(jnp.array([[1,0,0,0],[0,1,0,0],[0,0,1,1]]))
no_users, no_items = interactions.shape

def in_positives(user_id, item_id, interactions) -> bool:
  start, end = int(interactions.indptr[user_id]), int(interactions.indptr[user_id + 1])
  return item_id in interactions.indices[start:end]

print(in_positives(2, 3, interactions))

After reviewing multiple threads on the topic of dynamic slicing, I didn't find a solution. Is there any general approach to workaround such cases?

from functools import partial

import jax
import jax.experimental.sparse as jsp
import jax.numpy as jnp

interactions = jsp.BCSR.fromdense(jnp.array([[1,0,0,0],[0,1,0,0],[0,0,1,1]]))
no_users, no_items = interactions.shape

@partial(jax.jit, static_argnums=(0,1,))
def in_positives(user_id, item_id, interactions) -> bool:
  start, end = interactions.indptr[user_id], interactions.indptr[user_id + 1]
  return item_id in interactions.indices[start:end]

print(in_positives(2, 3, interactions))

IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>, Traced<ShapedArray(int32[])>with<DynamicJaxprTrace(level=1/0)>, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions).

[jakevdp]

Your code is attempting to create a dynamically-shaped array. JAX does not support dynamically-shaped arrays (see JAX Sharp Bits: Dynamic Shapes).

I don't know the full logic of the code you're trying to write, but I would suggest trying to express your logic without creating the dynamic array. You already have the statically-shaped source array and the dynamic start and end indices: can you write your remaining code to operate directly on those, and avoid the need for a dynamically-shaped intermediate array?

[dbalabka]

@jakevdp , thanks for such a quick reply.

In general, I would like to implement BPR-MF (Bayesian Personalized Ranking for Matrix Factorization) with JAX. The idea is to utilize JAX with XLA compilation to make it possible to utilize different architectures, GPUs/TPUs, and the power of multi-host /process computing. The original implementation is written in Cython, which, to my understanding, is hard to port even for GPU. So, inspired by JAX vs Cython benchmarks and previously presented Google ALS implementation in JAX (ALX: Large Scale Matrix Factorization on TPUs), I started to dig into similarly reimplementing BPR.

It took time to reimplement the method in vectorized form. So, I ended up with the following function implementation

from functools import partial
import timeit

import jax
import jax.experimental.sparse as jsp
import jax.numpy as jnp
from jax import Array

interactions = jsp.BCSR.fromdense(jnp.array([[1,0,0,0],[0,1,0,0],[0,0,1,1]]))
no_users, no_items = interactions.shape
interactions_indices_index = jnp.arange(0, interactions.indices.shape[0])
interactions_indices_mask = jnp.full_like(interactions.indices, -1)

@partial(jax.jit, static_argnums=(0,1,))
def in_positives(user_id, item_id, interactions) -> Array:
  start, end = interactions.indptr[user_id], interactions.indptr[user_id + 1]

  selected_indices = jnp.where(
    jnp.all(jnp.stack((start <= interactions_indices_index, interactions_indices_index <= end)), axis=0),
    interactions.indices,
    interactions_indices_mask
  )
  return jnp.any(selected_indices == item_id)

def in_positives_np(user_id, item_id, interactions) -> bool:
  start, end = int(interactions.indptr[user_id]), int(interactions.indptr[user_id + 1])
  return item_id in interactions.indices[start:end]

print(timeit.timeit(lambda: in_positives(2, 3, interactions), number=100))
print(timeit.timeit(lambda: in_positives_np(2, 3, interactions), number=100))

0.03485381699283607
0.16184099900419824

I should avoid using static arguments because static arguments are included in the cache key, and it leads to cache invalidation.

[jakevdp]

I don't think there's any need to mark those arguments as static here.

Also, take a look at https://jax.readthedocs.io/en/latest/faq.html#benchmarking-jax-code for tips on doing micro-benchmarks of JAX code – async dispatch and compilation overhead are probably in play here.

[dbalabka]

Thanks! Yes, definitely. I had read about these concepts, but I need to dig deeper to benchmark the code more precisely. I’m still very new to Jax

You are right about static arguments, and I managed to use traced values instead.