How to jit a function containing a 'drop diagonal' operation (or other mutations of arrays)

[xmax1]

If we want to create an operation which extracts the diagonal this can use the jax numpy api creating a function with 'non-static' array shapes

def function(x):
    x = x.reshape((5, 5))
    n = x.shape[0]
    y = jnp.diag(x)
    print(y.shape)
    return jnp.sum(y)

z = jax.jit(function)(x) # prints 5, returns no error

if we want to take all elements not the diagonal in numpy we might use a mask. This however doesn't work in jit because the array isn't static (this example of mutating arrays with masks is covered in the docs (https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html))

def function(x):
    x = x.reshape((5, 5))
    n = x.shape[0]
    eye = jnp.eye(n, dtype=bool)
    y = x[eye]
    return jnp.sum(y)

z = jax.jit(function)(x)

returns error: IndexError: Array boolean indices must be concrete.

is this kind of operation possible in Jax? Is any operation where we want to create a new tensor only consisting some elements of a previous tensor possible? Are there hacks to get around this?

[xmax1]

I think this function will work, but it's super inefficient

n = 4
square = np.random.normal(0., 1., (n, n))
def get_upper_triangle(square):
    n = square.shape[0]

    split1 = jnp.split(square, n, axis=0)
    split2 = [jnp.split(split1[i], [i], axis=1)[1] for i in range(0, n)]
    result = jnp.concatenate(split2, axis=1)
    print(square)
    print(result.shape)
    print(result)
    
jax.jit(get_upper_triangle)

get_upper_triangle(square)

[xmax1]

inefficient drop diagonal op

n = 4
square = np.random.normal(0., 1., (n, n))
def drop_diagonal(square):
    n = square.shape[0]
    split1 = jnp.split(square, n, axis=0)
    upper = [jnp.split(split1[i], [j], axis=1)[1] for i, j in zip(range(0, n), range(1, n))]
    lower = [jnp.split(split1[i], [j], axis=1)[0] for i, j in zip(range(0, n), range(0, n))]
    upper.extend(lower)
    result = jnp.concatenate(upper, axis=1)
    print(square)
    print(result.shape)
    print(result)
    
jax.jit(drop_diagonal)

drop_diagonal(square)