Concretize the value of an array within a JIT context

[kaushikcfd]

Is there a way to evaluate an expression within JIT context?

For example, in the snippet below could the variable doublify.result be computed eagerly to a DeviceArray?

def doublify(a):
    # Any way to eagerly evaluate *result*?
    # If not possible (for ex. "a" was dependent on placeholder arguments fed to the function being JIT-ed),
    # then raise.
    result = 2 * a
    return result

def get_array_of_twos():
    a = jnp.ones(10)
    return doublify(a)

jax.jit(get_array_of_twos)()

Thanks in advance!

[mattjj]

Great question! There is a way, though it's currently an internal API and we need to make it public (in the sense of: currently you have to import it from jax.core and we should move it to the top-level jax namespace, see #7535).

Well, there's almost a way. Your exact code as written doesn't work because by the time you call doublify it's already too late: the construction of the jnp.ones array has already been staged out as broadcast:

from functools import partial
import jax
import jax.numpy as jnp

@jax.named_call  # add this so it shows up in the jaxpr
def doublify(a):
  result = 2 * a
  return result

def get_array_of_twos():
  a = jnp.ones(10)
  return doublify(a)

print(jax.make_jaxpr(get_array_of_twos)())

{ lambda  ; .
  let a = broadcast_in_dim[ broadcast_dimensions=(  )
                            shape=(10,) ] 1.0
      b = named_call[ call_jaxpr={ lambda  ; a.
                                   let b = mul a 2.0
                                   in (b,) }
                      name=doublify ] a
  in (b,) }

The way to tell JAX, "make sure to evaluate this stuff at trace/compile time, or error if you can't!" is to use jax.core.eval_context, like this:

from functools import partial
import jax
import jax.numpy as jnp
import numpy as np

@jax.named_call  # add this so it shows up in the jaxpr
def doublify(a):
  with jax.core.eval_context()
    result = 2 * a
    print(f'result is {result}')
  return result

def get_array_of_twos():
  with jax.core.eval_context():
    a = jnp.ones(10)
  return doublify(a)

get_array_of_twos()

result is [2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]

Notice I had to use it twice: not only in doublify, but also where we create the constant jnp.ones. An alternative for that latter change would have been to import numpy as np and use np.ones there.

Maybe you can try using jax.core.eval_context, but leave a #TODO to use the public API endpoint (maybe jax.eval_context) when it's available? WDYT?

[kaushikcfd]

Thanks a lot for the reply!

Both the solutions seem a bit sub-optimal for the use case I'm targeting. My use case is to evaluate a sub-expression returned by a function and store that into DeviceArrays. The catch is the function in this case has its result memoized and it is called from both JIT and non-JIT contexts so a jax.named_call won't quite work, and, wrapping exactly the subexpression into a jax.core.eval_context is tedious (and impractical) to integrate into the existing code base.

I'm assuming since the data structure representing the computation must be a DAG, couldn't we do a topological sort traversal of the graph and evaluate the result at the end?

(Sorry had added this as answer, moved it as a reply)