Jax with reversible-jump MCMC

[mhinne]

I'm trying to implement a specific reversible-jump MCMC algorithm, By definition, an RJ algorithm traverses solutions of different dimensionalities, which is obviously tricky with Jax, but I'm hoping I'm simply missing a clever trick to make this work :-)

Here is a MWE of what I am trying to do:

import jax.numpy as jnp
import jax.random as jrnd

key = jrnd.PRNGKey(0)
key, wishart_key, choice_key = jrnd.split(key, 3)
p = 4
dof = p + 1
scale = jnp.eye(p)

# Create a pxp Wishart-distributed random matrix
u = jrnd.multivariate_normal(wishart_key, mean=jnp.zeros((p, )), cov=scale, shape=(dof, ))
Sigma = jnp.einsum('dp,dq->pq', u, u)

# Create a pxp symmetric, binary matrix
G = jnp.array([[1, 1, 0, 0], [1, 1, 1, 0], [0, 1, 1, 1], [0, 0, 1, 1]])  


def take_dynamic_submatrices(i, G):
    # select all variables that i connects to in G
    ix = jnp.where(G[i] == 1)[0]
    W_sub = W[ix[:, None], ix]
    return W_sub

#

# Pick a node
i = 1

W_sub = take_dynamic_submatrices(i, G)
print('W_sub:\n', W_sub)

# How can we make this work within a vmap or lax.scan?
W_sub_jit = jax.jit(take_dynamic_submatrices)(i, G)
print('W_sub_jit:\n', W_sub_jit)

Throughout the algorithm, Sigma/W and G would change values, but remain of shape $p\times p$. However depending on random steps, G would change, and that would mean ix and hence W_sub would/could be different in size across iterations of the algorithm. I do not really need to store elements of different sizes, as I would assign W_sub back to a subpart of W, keeping what I track over iterations of the same shape. I would need to have them as variables on the fly though.

I have tried a couple of directions: static_argnums, and simply masking the matrices that I update. The first one does not work, because G is not hashable (as it is an array), and the second approach failed as well because later I have to do some linear algebra operations on W_sub, and that does not work if it is masked.

It's definitely possible that what I am trying to do is simply impossible with Jax due to the requirement of array shapes being known, but since the objects I really care about do keep the same shape, do you perhaps have a suggestion on how to proceed?

[jaro-sevcik]

For JIT to work, even temporary variables, such as W_sub, must have static shapes. As a result, what you ask is not possible unless you make the shapes static. However, there are other approaches than masking (which can indeed get messy) - if you were fine with zero-padding W_sub, that is possible from JITted code:

import jax.numpy as jnp
import jax.random as jrnd

key = jrnd.PRNGKey(0)
key, wishart_key, choice_key = jrnd.split(key, 3)
p = 4
dof = p + 1
scale = jnp.eye(p)

# Create a pxp Wishart-distributed random matrix
u = jrnd.multivariate_normal(wishart_key, mean=jnp.zeros((p, )), cov=scale, shape=(dof, ))
Sigma = jnp.einsum('dp,dq->pq', u, u)
W = Sigma

# Create a pxp symmetric, binary matrix
G = jnp.array([[1, 1, 0, 0], [1, 1, 0, 1], [0, 1, 1, 1], [0, 0, 1, 1]])  

def take_dynamic_submatrices(i, G):
    # select all variables that i connects to in G
    ix = jnp.where(G[i] == 1)[0]
    W_sub = W[ix[:, None], ix]
    return W_sub

def take_dynamic_submatrices_padded(i, G):
    # select all variables that i connects to in G
    mix = jnp.arange(len(G[i]))
    mix = jax.lax.associative_scan(lambda a, b: a + b, G[i]) - 1
    size = mix[-1] + 1
    mix = jnp.where(G[i] == 1, mix, len(G[i]) + 1)
    W_sub = jnp.zeros(W.shape).at[mix[:, None], mix].set(W, mode='fill')
    return W_sub, size

# Pick a node
i = 1

W_sub = take_dynamic_submatrices(i, G)
print('W_sub:\n', W_sub)

# How can we make this work within a vmap or lax.scan?
W_sub_jit, size = jax.jit(take_dynamic_submatrices_padded)(i, G)
print('W_sub_jit:\n', W_sub_jit)
print('W_sub_jit size:\n', size)

Output:

W_sub:
 [[ 5.8744636 -0.868772  -1.575294 ]
 [-0.868772   6.0721636  2.6670828]
 [-1.575294   2.6670828  2.2540443]]
W_sub_jit:
 [[ 5.8744636 -0.868772  -1.575294   0.       ]
 [-0.868772   6.0721636  2.6670828  0.       ]
 [-1.575294   2.6670828  2.2540443  0.       ]
 [ 0.         0.         0.         0.       ]]
W_sub_jit size:
 3

[jakevdp]

You cannot do what you're asking: i.e. create a matrix W_sub that has a dynamic shape.

In general you probably can do something similar to what you have in mind, by side-stepping the actual creation of the dynamic matrix W_sub. But your MWE is not enough to offer advice on that, because the best approach will depend on what operations you're doing with W_sub.

For example:

• if you're doing an aggregation over W_sub, then rather than W_sub = W[ix[:, None], ix] you could do W_sub = jnp.where(ix[:, None] & ix], W_sub, jnp.nan) and then use a NaN-aware version of the aggregation
• if you're doing matrix multiplication with W_sub, you could use W_sub = jnp.where(ix[:, None] & ix], W_sub, 0) and then multiply with a similarly constructed vector
• if you're performing an SVD on W_sub then you could use ix = jnp.where(G[i] == 1, size=len(G), fill_value=0)[0] and then perform the SVD on the resulting W_sub, and only work with the first len(G) elements of the singular vectors.

If you need specific advice, you could expand your MWE to show the kinds of operations you're trying to do.

[mhinne]

Thank you very much! With W_sub I primarily need to do a Cholesky decomposition (and some solves). Would that work in a similar way to your suggestion on SVD?