detect subdiagonal nonzero entries

[hyperfra]

Hello
How I need to detect the subdiagonal non zero entries of quasi-triangular matrix R (which is computed several times)
In python I have:
r = np.where(np.abs(np.diag(R, -1)) > 100*eps)[0]

with Jax the issue is that the number of subdiagonal elements that are non zeros can change (not the overall size of the matrix). Any smart way of doing that?

[jakevdp]

If you want to know the number of nonzero entries, you can use:

np.count_nonzero(np.abs(np.diag(R, -1)) > 100*eps)

If you want to create an array with only the nonzero entries, that's not possible within a JAX transform because the number of entries is dynamic. Depending on what you want to do with those entries, there may be a number of solutions: for example you could create a padded the array to a static size by using the size argument of jnp.where.

[hyperfra]

Thank you Jake
I think I understand the logic. However the size argument does not appear to solve the issue. I keep the size n and use it to create an array with the indexes of the rows:
n = jnp.count_nonzero(jnp.abs(jnp.diag(R, -1)) > 100*epsilon)
r = jnp.where(jnp.abs(jnp.diag(R, -1)) > 100*epsilon, size=n)[0]
still gives me a ConcretizationTypeError.

[jakevdp]

Understood: the point of the size argument is to specify a static array size. n here is not static (it is computed from the contents of an input array), so this is not a solution to your problem. Just to emphasize: there's no way to do precisely what you are asking, because it requires generating a dynamcially-sized array, and JAX cannot generate dynamically-sized arrays within transformations. You'll need to change your approach to use a non-dynamic arrays size instead.

For example, you could pass size=R.shape[0] - 1, because array shapes are always static. The result would be a statically sized array with the first n elements set to the values you are interested in, and the rest filled with 0. Depending on what you hope to do with this output, this may or may not be a suitable approach.

[hyperfra]

Thank u. I was trying to jit an algorithm to sort the Schur decomposition according to the eigenvalues, and this requires to know where the matrix has 2x2 blocks in the diagonal. Would it be a feasible strategy to explode the algorithm in many functions so that eventually the dynamic shape becomes static input of another function? or it is hopeless.