Inequation inside Jax.jit function

[fasghq]

Hi!
#2508 (comment)
Maybe I miss something obvious, but what should I do if I need to use inequation inside the function for jax.jit? Here's my function:
def tt_round_model(tt, delta):
...
for x in s:
if x < delta:
...
As I found, I should add static_argnums in such case, so I added delta to it:
fast_tt_round = jit(tt_round_model, static_argnums=(1))
ttr = fast_tt_round(tts_s, eps / math.sqrt(len(tts_s.core) - 1))
But I still have the problem in compilation:

I suppose that the problem is in my x variable which is diagonal element of s matrix which I get inside this algorithm, so I can't add it to static_argnums.
Thank you for your help!

[jakevdp]

In general, you cannot jit-compile code with control flow that depends on the value in an array. It's possible to use structured control flow primitives in this case, but this sort of loop-based construction will generally be much less efficient than array-based primitives, particularly on GPU or TPU.

For code like yours, it would likely be best to write your code using broadcasting and jnp.where:

result = jnp.where(s < delta, value1, value2)

[fasghq]

When I used jnp.where I got

[jakevdp]

Can you give a more complete example of what you're trying to do? I was guessing at what might work given the handful of lines you included in your question.

[fasghq]

I work with Tensor-Train and this function rounds tt (decreases its ranks). The function takes tt class with the list of cores on input:

@struct.dataclass 
class Model():
    core: list

As I need to compare singular values with delta, I added delta to input and put it to the static_argnums. In the second for-loop, I calculate svd of k-th core G[k] receiving matrix s with singular values, then I need to cut singular values which are less than delta.

def tt_round_model(tt, delta):  
    G = tt.core
    d = len(tt.core)
    rn = []
    time1 = time.perf_counter()    
    
    for k in range(d-1, 0, -1):
        r1, n, r2 = G[k].shape
        G[k] = np.reshape(G[k], (r1, n * r2))
        G[k] = np.transpose(G[k])
        G[k], R = la.qr(G[k])
        G[k] = np.transpose(G[k])
        G[k-1] = np.einsum('ijk,lk->ijl', G[k-1], R) 
    
    for k in range(0, d - 1):
        print(k, G[k].shape, k+1, G[k+1].shape)
        
        if k == 0:
            r1, n1, r2 = G[k].shape
            r3 = G[k+2].shape[0]
            n2 = int(G[k+1].shape[1] / r3)
        elif k == d - 2:
            r1 = G[k].shape[0]
            r2 = G[k+1].shape[0]
            n1 = int(G[k].shape[1] / r2)
            r3 = 1
            n2 = G[k+1].shape[1]
        else:
            r1 = G[k].shape[0]
            r2 = G[k+1].shape[0]
            n1 = int(G[k].shape[1] / r2)
            r3 = G[k+2].shape[0]
            n2 = int(G[k+1].shape[1] / r3)

        G[k] = np.reshape(G[k], (r1 * n1, r2))
        G[k+1] = np.reshape(G[k+1], (r2 , n2 * r3))
        u, s, v = la.svd(G[k])

        rc = s.shape[0]
        p = 0
        
        '''
        s = jnp.where(s > delta)
        rc = s.shape[0]
        '''
        for x in s:
            if x < delta:
                x = 0
                if rc > p:
                    rc = p
            p = p + 1
        
        s = np.diag(s)
        u = u[:, :rc]
        s = s[:rc, :rc]
        v = v[:, :rc]
        v = np.dot(v, s)
        G[k] = np.reshape(u, (r1, n1, rc))   
        G[k+1] = np.einsum('ij,ik->jk', G[k+1], v)
    
    
    if len(G[d-1].shape) == 2:
        G[d-1] = np.reshape(G[d-1], (G[d-1].shape[1], G[d-1].shape[0], 1))

    for X in G:
        rn.append(X.shape[0])
    rn.append(G[d-1].shape[2])

    time2 = time.perf_counter()
    print("\n time: ", time2 - time1)
    
    return Model(G)

At the output of the function, I get ttr (rounded tt) with decreased ranks of cores:

fast_tt_round = jit(tt_round_model, static_argnums=(1))
ttr = fast_tt_round(tts_s, eps / math.sqrt(len(tts_s.core) - 1))

[jakevdp]

What is x=0 meant to do within your loop? It looks like a no-op, because the variable is never used.

It looks like you can replace this block

for x in s:
    if x < delta:
        x = 0
        if rc > p:
            rc = p
    p = p + 1

with

p = len(s)
rc = max(rc, (s < delta).sum())

which will generally be much more efficient, and will have no issue with compilation.

[jakevdp]

Yes, that should work. But be aware that using a for loop in that way can lead to very long compile times and likely slower runtimes. You'll see much better performance if you use array-based operations directly.

[fasghq]

As I understand I can't use jax.dynamic_slice here. I found a solution to a similar problem:

@jax.jit
def sum_first_k(a, k):
  n = len(a)
  return np.sum(np.where(np.arange(n) < k, a, 0))

I suppose that something like:

n1 = u.shape[0]
n2 = u.shape[1]
k1 = u.shape[0]
k2 = rc
u = np.where([np.arange(n1) < k1, np.arange(n2) < rc], u, ???)

should work as u = u[:, :rc], but what should I put instead ??? to not to put extra elements to u (when condition is False)?

[jakevdp]

I'm sorry, it's not clear to me what you are trying to accomplish in that example.

[fasghq]

I mean using

u = u[:, :rc]
s = s[:rc, :rc]
v = v[:, :rc]

causes error

IndexError: Array slice indices must have static start/stop/step to be used with Numpy indexing syntax. Try lax.dynamic_slice/dynamic_update_slice instead.

But I can't use lax.dynamic_slice here as

n1 = u.shape[0]
u = lax.dynamic_slice(u, (0, 0), (n1, rc))

causes error

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected (in `bool`).
Use transformation parameters such as `static_argnums` for `jit` to avoid tracing input values.
See `https://jax.readthedocs.io/en/latest/faq.html#abstract-tracer-value-encountered-where-concrete-value-is-expected-error`.
Encountered value: Traced<ShapedArray(bool[]):JaxprTrace(level=-1/1)>

[fasghq]

Probably I should better create new issue.