Functions of Index Arrays -- Efficient Calculation

[danielkelshaw]

Hi, some I'm currently doing some work to compute Christoffel symbols using jax, essentially of the form:

$$ \Gamma_{kij} = \frac{1}{2} \left( \frac{\partial g_{ki}}{\partial x^{j}} + \frac{\partial g_{kj}}{\partial x^{I}} - \frac{\partial g_{ij}}{\partial x^{k}} \right) $$

I have a function to compute the metric tensor $g_{ij}$ and can compute the required jacobin using jax.jacobian just fine.

Constructing $\Gamma_{kij}$
I would like to find a way by which to construct the $\Gamma_{kij}$ without nested for loops. As the dimensionality of the system increases, the nested for loop becomes impractical.

My first attempt at solving this was simple nested for loops:

def fk_christoffel(x: jnp.ndarray, metric: Callable[[jnp.ndarray], jnp.ndarray]) -> jnp.ndarray:
        
    dgdx = jax.jacobian(metric)(x)
    
    christ_fk = jnp.zeros_like(dgdx)
    for k, i, j in it.product(*map(range, dgdx.shape)):
        christ_fk = christ_fk.at[k, i, j].set(dgdx[k, i, j] + dgdx[k, j, i] - dgdx[i, j, k])
        
    christ_fk /= 2.0
    
    return christ_fk

I considered writing this explicitly with jax.lax.fori_loop but this doesn't really solve the problem.

I then considered first generating all of the possible permutations and using jax.vmap to update the tensor. The permutations have to be computed only once for a given input shape so I make use of ft.lru_cache here:

@ft.lru_cache(maxsize=16)
def index_permutations_of_shape(shape: tuple[int]) -> list[tuple[int]]:
    return jnp.array([p for p in it.product(*map(range, shape))])
        
def fk_christoffel(x: jnp.ndarray, metric: Callable[[jnp.ndarray], jnp.ndarray]) -> jnp.ndarray:
      
    dgdx = jax.jacobian(metric)(x)
        
    def set_value(kij: jnp.ndarray, _christ_fk: jnp.ndarray) -> jnp.ndarray:
        k, i, j = kij
        return _christ_fk.at[k, i, j].set(dgdx[k, i, j] + dgdx[k, j, i] - dgdx[i, j, k])
    
    christ_fk = jnp.zeros_like(dgdx)
    
    perms = index_permutations_of_shape(dgdx.shape)
    christ_fk = jax.vmap(set_value, in_axes=(0, None))(perms, christ_fk).sum(axis=0)
        
    christ_fk /= 2.0
    
    return christ_fk

While this seems to work okay, the memory requirements for large shapes will be huge - does jax handle this internally?

I would like to use something like jnp.fromfunction to do this but cannot get it to work with array indexing. Are there any good alternative approaches which would avoid the combinatorial explosion with increasing dimension?

[danielkelshaw]

I believe I have found an efficient implementation which uses jnp.vectorize:

def fk_christoffel(x: jnp.ndarray, metric: Callable[[jnp.ndarray], jnp.ndarray]) -> jnp.ndarray:
      
    dgdx = jax.jacobian(metric)(x)
    
    def get_value(k, i, j) -> jnp.ndarray:
        return 0.5 * (dgdx[k, i, j] + dgdx[k, j, i] - dgdx[i, j, k])

    return jnp.vectorize(get_value)(*jnp.indices(dgdx.shape))

I would appreciate any feedback on if there's a better way to do this still -- if not, please feel free to close.

Thank you!

[mattjj]

This is so cool!

On a first skim I was going to suggest jnp.vectorize, which is a wrapper around jax.vmap (i.e. it just calls jax.vmap multiple times). But indeed you may run out of memory. If you do, xmap may be a good fit, since in addition to being able to vmap over multiple axes (like jnp.vectorize) it also lets you tune how much to vectorize vs how much to use as a sequential loop. You can also parallelize over devices.

[danielkelshaw]

xmap looks like a good option, thanks for the recommendation -- I'll take a proper look and see if I can get something working.