Applying a sequence of Householder transformations with jax.lax-style loops

[spenrich]

Hello! I'm working on a function that maps parameters to an orthogonal matrix. This involves a sequence of Householder transformations, where only a sub-block of the matrix iterate is operated on at each iteration. Below is a working version with a regular Python for-loop.

from math import isqrt

import jax
import jax.numpy as jnp


def dim_to_tri_size(n):
    """Compute the number of elements in the lower-triangle of an array."""
    d = n * (n + 1) // 2
    return d


def tri_size_to_dim(d):
    """Compute the square matrix dimension with `d` lower-triangle elements."""
    n = (isqrt(8*d + 1) - 1) // 2
    return n


def params_to_ortho(θ):
    """Map parameters to a unique special-orthogonal matrix.
    
    `V` is assumed to be a square upper-triangular matrix of dimension `n - 1`,
    where the last `n - i` elements of row `i` contain the vector defining
    the Householder matrix `H(i)`. The output `Q` is given by:
        `Q = H(1) @ H(2) @ ... @ H(n)`.
    """
    θ = jnp.asarray(θ)
    d = θ.shape[-1]
    n = tri_size_to_dim(d) + 1

    # Initialize Q = H(n). Normally any Householder matrix `H` satisfies 
    # `det(H) = -1`, so to ensure `det(Q) = 1` we set `H(n) = I` when `n` is 
    # odd. Otherwise, `H(n)` looks like `I` with a `-1` in the bottom-right.
    Q = jnp.eye(n).at[-1, -1].set(
        jnp.where(n % 2 == 0, -1., 1.)
    )

    # Iteratively pre-multiply Householder reflectors, such that
    #   `Q = H(1) @ H(2) @ ... @ H(n)`.
    # None of the Householder matrices are explicitly constructed. Since `H(i)`
    # only affects the bottom-right conformable sub-matrix of `Q`, we carry out
    # the required matrix-vector operations on only this sub-matrix at each
    # iteration.
    one = jnp.ones((1,))
    k = d
    for i in range(1, n):  # TODO: use LAX loop
        v = θ[k-i:k]
        τ = 2. / (1. + jnp.sum(v**2))
        v1 = jnp.concatenate((one, v))
        p = i + 1  # size of sub-matrix
        Q = Q.at[-p:, -p:].add(-jnp.outer(τ*v1, Q[-p:, -p:].T @ v1))
        k -= i
    return Q

Each iteration of the for-loop involves an operation on a conformable bottom-right sub-block of the carry value Q. I want to replace the for-loop with a jax.lax-style loop. Below is my attempt with jax.lax.fori_loop.

def params_to_ortho_lax(θ):
   """Map parameters to a unique special-orthogonal matrix.
   
   `V` is assumed to be a square upper-triangular matrix of dimension `n - 1`,
   where the last `n - i` elements of row `i` contain the vector defining
   the Householder matrix `H(i)`. The output `Q` is given by:
       `Q = H(1) @ H(2) @ ... @ H(n)`.
   """
   θ = jnp.asarray(θ)
   d = θ.shape[-1]
   n = tri_size_to_dim(d) + 1

   # Initialize Q = H(n). Normally any Householder matrix `H` satisfies 
   # `det(H) = -1`, so to ensure `det(Q) = 1` we set `H(n) = I` when `n` is 
   # odd. Otherwise, `H(n)` looks like `I` with a `-1` in the bottom-right.
   Q = jnp.eye(n).at[-1, -1].set(
       jnp.where(n % 2 == 0, -1., 1.)
   )

   # Iteratively pre-multiply Householder reflectors, such that
   #   `Q = H(1) @ H(2) @ ... @ H(n)`.
   # None of the Householder matrices are explicitly constructed. Since `H(i)`
   # only affects the bottom-right conformable sub-matrix of `Q`, we carry out
   # the required matrix-vector operations on only this sub-matrix at each
   # iteration.
   def body(i, carry):
       k, θ, Q = carry
       n = Q.shape[-1]
       one = jnp.ones((1,))
       v = θ[k-i:k]
       τ = 2. / (1. + jnp.sum(v**2))
       v1 = jnp.concatenate((one, v))
       p = i + 1  # size of sub-matrix
       Q = Q.at[-p:, -p:].add(-jnp.outer(τ*v1, Q[-p:, -p:].T @ v1))
       k -= i
       carry = (k, θ, Q)
       return carry

   carry_init = (d, θ, Q)
   carry = jax.lax.fori_loop(1, n, body, carry_init)
   _, _, Q = carry
   return Q

However, the test code below

from tqdm.auto import tqdm

# Configure JAX and set random seed
jax.config.update("jax_enable_x64", True)
seed = 0
key = jax.random.PRNGKey(seed)

n_max = 5
for n in tqdm(range(2, n_max + 1)):
    d = dim_to_tri_size(n) - 1
    θ = jax.random.uniform(key, (d,))
    I = jnp.eye(n)
    Q = params_to_ortho(θ)
    Q_lax = params_to_ortho_lax(θ)
    assert jnp.allclose(Q@Q.T, I)
    assert jnp.allclose(Q.T@Q, I)
    assert jnp.allclose(jnp.linalg.det(Q), 1.)
    assert jnp.allclose(Q, Q_lax)

gives me the error:

Traceback (most recent call last):
  File "/home/spenrich/orthogonal.py", line 113, in <module>
    Q_lax = params_to_ortho_lax(θ)
  File "/home/spenrich/orthogonal.py", line 97, in params_to_ortho_lax
    carry = jax.lax.fori_loop(1, n, body, carry_init)
  File "/home/spenrich/orthogonal.py", line 87, in body
    v = θ[k-i:k]
  File "/home/spenrich/opt/anaconda3/envs/asl/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 3828, in _rewriting_take
    return _gather(arr, treedef, static_idx, dynamic_idx, indices_are_sorted,
  File "/home/spenrich/opt/anaconda3/envs/asl/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 3837, in _gather
    indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
  File "/home/spenrich/opt/anaconda3/envs/asl/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 4089, in _index_to_gather
    raise IndexError(msg)
IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(Traced<ShapedArray(int64[], weak_type=True)>with<DynamicJaxprTrace(level=1/0)>, Traced<ShapedArray(int64[], weak_type=True)>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).

It seems that the jax.lax-style loop does not like how the size of the vector v is dynamic with the way I specified it. However, I know beforehand how large v will be for each iteration, so I wonder if there is some way I can implement this without dynamic slicing. Ultimately, I'd like to get this working with jax.lax.scan in particular. Any idea how I can fix this?

[jakevdp]

You can not do jax.lax-style loops over arrays with varying size. The only way to do this would be to somehow modify your code such that it uses fixed-size arrays, perhaps through some form of manual masking.

[patrick-kidger]

As Jake says, a lax-style loop over variable-size arrays doesn't work.

One option might be to see if you can express your variable-size computation using a lax.while_loop. That is, do lax.fori_loop(f, ...) where

def f(i, carry):
    ...
    lax.while_loop(...)
    ...

the body function has another loop inside of it.