create a concrete value during the tracing

[mariogeiger]

Hi,

Is there a way to create concrete arrays while being in a tracer call?

@jit
def f():
    x = jnp.array([1.0, 2.0])
    # I want a concrete x because I want to use it only during the tracing

Here is why I need that. Let say I have two cases whether a function is symmetric or antisymmetric and I have a function to check that numerically:

def parity(f):
    a = jnp.linspace(-1.0, 1.0, 21)
    if jnp.allclose(f(a), f(-a)):
        return 'symmetric'
    if jnp.allclose(f(a), -f(-a)):
        return 'antisymmetric'
    return 'undefined'

The problem is that if I use parity inside a function that gets jited I get an error because a is not concrete:

ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 

How to force a to be concrete?

[mariogeiger]

I dream of something like a context manager

def parity(f):
    with jax.concrete():
        a = jnp.linspace(-1.0, 1.0, 21)
        if jnp.allclose(f(a), f(-a)):
            return 'symmetric'
        if jnp.allclose(f(a), -f(-a)):
            return 'antisymmetric'
        return 'undefined'

[jakevdp]

The easiest way to make a value (or an operation) concrete, is to use the np version rather than the jnp version:

def parity(f):
    a = np.linspace(-1.0, 1.0, 21)
    if np.allclose(f(a), f(-a)):
        return 'symmetric'
    if np.allclose(f(a), -f(-a)):
        return 'antisymmetric'
    return 'undefined'

[mariogeiger]

Amazing !

[jakevdp]

I don't think core.eval_context is considered a public API, so if you use it be prepared for it to change without warning in future releases.

[mariogeiger]

Of course

[bonfab]

I have similar problem where this would also be the most elegant solution. Do plans exist to make a jax context manager officially available?

[soraros]

FYI, there is #7535