Error while using vmap to iterate over an array to provide value for the 'n' parameter of jax.numpy.linalg.matrix_power function

[rkmalaiya]

Hi All,

I am trying to implement the matrix_power function for 'n' participants. The matrix_power function raises a square matrix to a whole number. e.g., it will raise the MxM matrix to 5.

I have an MxM matrix for N participants. So, I have a matrix of NxMxM. I need to iterate over the 0th dim (N) and raise the MxM matrix to a whole number which is different for each participant. I assumed vmap might be a great way to do this, but I keep getting below error. The error is because, the matrix_power function expects a scaler but instead is getting a Traced datatype.

I thought about using a simple for loop, but I am building a likelihood function that I will then give to numpyro for MCMC inference. Hence, I thought that vmap may better facilitate autodiff and inference performance. I could be wrong, please advise if there's a way to fix the below error or should I just try for-loop?

Code:

from jax import numpy as jnp
from jax import vmap

N = 2
M = 3
T_t = jnp.ones((N,M,M)) 
n = jnp.arange(1,N+1) 
def mat_pwr (T_t, n): 
    return jnp.linalg.matrix_power(T_t, n)

ans = vmap(mat_pwr, in_axes=(0, 0))(T_t, n)

Error:

TracerIntegerConversionError: The index() method was called on traced array with shape int32[].
...
TypeError: exponent must be an integer, got Traced<ShapedArray(int32[])>with<BatchTrace(level=1/0)> with
val = Array([1, 2], dtype=int32)
batch_dim = 0

The code works if I just execute the matrix_power function for a single value of 'power' parameter n

Code:

i=0
jnp.linalg.matrix_power(T_t[i,...], n[i].item())

Also, below code works if I just manually select an item from the 'n' array. But, I need to iterate through all the values of 'n'.

from jax import numpy as jnp
from jax import vmap

N = 2
M = 3
T_t = jnp.ones((N,M,M)) 
n = jnp.arange(1,N+1) 
def mat_pwr (T_t): 
    return jnp.linalg.matrix_power(T_t, n[0].item())

ans = vmap(mat_pwr, in_axes=(0))(T_t)

Thanks for the help!

[jakevdp]

Hi - thanks for the question. Currently the n argument to matrix_power must be a static integer, meaning that it's not possible to vmap over this argument, because mapped arguments are not static. In that sense, this is working as expected given the current implementation of jnp.linalg.matrix_power.

Unless that changes, the optimal approach here is probably to avoid the vmap and use a Python for-loop instead; i.e. something like this:

from jax import numpy as jnp
from jax import vmap

N = 2
M = 3
T_t = jnp.ones((N,M,M)) 
n = jnp.arange(1,N+1) 

result = jnp.stack([jnp.linalg.matrix_power(mat, int(pow))
                    for mat, pow in zip(T_t, n)])
print(result)

[[[1. 1. 1.]
  [1. 1. 1.]
  [1. 1. 1.]]

 [[3. 3. 3.]
  [3. 3. 3.]
  [3. 3. 3.]]]

[jakevdp]

One other note: if you're dealing only with hermetian matrices, there is a straightforward way to write a closed-form matrix power:

def matrix_power_h(A, n):
  evals, evecs = jnp.linalg.eig(A)
  return evecs @ jnp.diag(evals**n) @ evecs.conj().T

This doesn't rely on n being static, and so it would work with your vmap solution, but would not return the correct matrix power if the input is not actually Hermitian.

[rkmalaiya]

Thanks a lot for the response. It's not a Hermitian matrix. So I will use the for-loop solution. I am assuming the auto differentiation will work with the for-loop.

[jakevdp]

Yes, autodiff will work with a for-loop.

[rkmalaiya]

Hi All,

I am sorry for posting this question here given this discussion is already marked as answered by me. But this question is related to this thread so I thought may be I'll ask the question here.

I went ahead and implemented the for-loop for matrix_power function as discussed with @jakevdp above (#18356 (comment)). The matrix_power function is part of a likelihood function that I am estimating using NUTS method in NumPyro. But, I am facing difficulty in the compilation phase of posterior sampling. There are few other operations in the likelihood function but they are all simple dot products. I have also mentioned the full equation below. For just N=10, M=11 the compilation phase took more than 24 hrs and then I had to kill the job. So I don't know if it would have even completed. For N=10 and M=3, it compiles in 20-30 mins. My requirement is N=400 and M=11. So, is there a way I can parallelize/optimize this code for faster compilation. More importantly, is this compilation time due to the for-loop that we had to use?

A @ [ B @ T ]^N @ p
here,
A is MxM
B is MxM
T is NxJxMxM
N is NxJ
p is Mx1
^ denotes matrix_power function.

Thanks for all the help!