JAX doesn't need math.comb to infer array shapes?

[ajipalar]

I wrote a custom function to get the lower triangular values of a matrix (I am aware of A[jnp.tril_indices(n, k=-1)])
The function depends on a call to math.comb. Somehow JAX seems to correctly trace the function without this call. How does this work?

def matrix2flat(M, row_major=True):
    """
    An iterative method to flatten a matrix
    """
    n, m = M.shape
    N = math.comb(n, 2) # questionable line
    a = jnp.zeros(N, dtype=M.dtype)
    k=0
    if row_major:
        for row in range(n):
            l = m - row - 1
            a = a.at[k:k + l].set(M[row, row + 1:m])
            k += l
    else:
        for col in range(m):
            l = n - col - 1
            a = a.at[k:k + l].set(M[col+1:n, col])
            k += l
    return a

For a specific case of a (5, 5) matrix and an expected output of shape (10,), inspecting the first few lines of the JAXPR shows

{ lambda ; a:f32[5,5]. let
    b:f32[10] = broadcast_in_dim[broadcast_dimensions=() shape=(10,)] 0.0
    c:i32[1] = broadcast_in_dim[broadcast_dimensions=() shape=(1,)] 0
    ...

How does JAX know that shape of the flattened array should be (10,)? I know from previous testing that JAX fails to trace
math.comb.

My only guess is that during tracing JAX is getting the array size from the final loop index l and setting the shape to (l+1,).
If this is what is happening it seems like black magic because the shape of a should have been defined before dropping into the loop.

JAX Version: 0.4.8

[jakevdp]

In JAX, array shapes are static and represented by Python integers. Any operation on these static integers will be done at trace-time in Python: so math.comb(n, 2) is just a Python function computed between two Python integers. The result of any such static argument will be embedded as a constant in the XLA program that is eventually compiled.

For more on JAX's computational/tracing model, you might find this doc useful: https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html

[ajipalar]

Partial application of math.comb(N, 2) works. Didn't realize operations on static ints were done at trace-time. Thanks for the quick reply!