A jit puzzle on how to tell that the size of arrays are actually known in the context of a set diff and union problem

[shailesh1729]

Consider the following simplified example

import jax.numpy as jnp
import jax
n = 16
k = 3
# A given set of k unique indices
A = jnp.array([7, 2, 5], dtype=jnp.int32)
# Another set of 2*k unique indices 
B = jnp.array([4, 3, 2, 1, 9, 10], dtype=jnp.int32)


# Find k new indices in B and combine them with A
def extend_A_by_k_from_B(A, B, k):
    C = jnp.setdiff1d(B, A)
    D = jnp.hstack((A, C[:k]))
    return D

D = extend_A_by_k_from_B(A, B, k)
print(D)


f = jax.jit(extend_A_by_k_from_B, static_argnums=(2,))
D = f(A, B, k)
print(D)

In the function extend_A_by_k_from_B:

• A and B are both index arrays for some (unspecified) larger array.
• Size of A is known to be k
• Size of B is known to be 2*k
• It is known that both A and B have unique indices individually, i.e., no duplicates within A or B.
• Intersection of A and B cannot have more than "k" common indices
• It is known that B has at least "k" new indices, not in A.
• The result D has to pick k indices from A and new k indices from B.
• Thus, the size of D is known in advance to be 2*k by the way the function has been designed.
• Thus, the size, shape, and type of input arrays, and output array are known in advance.

While the non-JIT version works fine, I cannot figure out a way to convert this logic into a form that is compatible with JIT compilation. The error I face is:

jax._src.traceback_util.UnfilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(int32[5])>with<DynamicJaxprTrace(level=0/1)>

I tried various ideas using arrays of flags with jnp.where type logic but could not make this logic work with JIT.

Is there a way to rewrite this logic that can be happily accepted by the JIT compiler?

P.S.: This is a simplified set union step in an algorithm called Compressive Sampling Matching Pursuit.

[jakevdp]

Here's a version that might work for you, given the assumptions about the inputs

def extend_A_by_k_from_B(A, B, k):
  ind = jnp.argsort(jnp.in1d(B, A))[:k]
  return jnp.hstack((A, B[ind]))

It chooses the first k inputs of B that do not appear in A, and concatenates them to A