Problem in porting scipy.linalg.orth with JIT

[shailesh1729]

A straightforward port of scipy.linalg.orth is the following.

import jax.numpy as jnp
from jax import jit, lax
from jax.scipy.linalg import svd

def orth(A, rcond=None):
    u, s, vh = svd(A, full_matrices=False)
    M, N = u.shape[0], vh.shape[1]
    if rcond is None:
        rcond = jnp.finfo(s.dtype).eps * max(M, N)
    tol = jnp.amax(s) * rcond
    num = jnp.sum(s > tol, dtype=int)
    Q = u[:, :num]
    # Q = lax.dynamic_slice(u, (0, 0), (M, num))
    return Q

# the following doesn't work
orth_jit  = jit(orth, static_argnums=(1,))

The key problem is that number of columns in the Q is data-dependent and JAX doesn't support dynamic slices. Even lax.dynamic_slice doesn't help since the variable num is also a traced object and cannot be converted to an integer.

Is this the reason why it is not included in jax.scipy.linalg package? Is there some trick that can alleviate the JIT compilation issue for this function?

One possible workaround I was thinking of was to change the return type to a tuple (u, num). Would that be the right approach?

[hawkinsp]

Yes, in general JAX does not (yet) support data-dependent shapes inside a jit. Here, num necessarily depends on the matrix. Outside a jit it works fine.

The usual workaround in JAX is to return a padded array, i.e., return Q unsliced together with num that tells you how much of the array contains useful data. This requires an API change; sometimes we do that by adding an optional keyword argument that states if you want to exactly mimic the NumPy/SciPy behavior or return the (Q, num) pair in a jit-compatible way.

We have plan for more first-class support for dynamic shapes inside jit but nothing ready at this time.

[shailesh1729]

Thanks. It's good to know that dynamic shapes are indeed on the roadmap for jit. Will definitely be waiting for the first-class support of dynamic shapes in jit.