No crash in JIT? I was expecting one....

[jecampagne]

Hi,
Although astonishing sometimes I am curious to know why a code is not crashing :)

In the thread of trace-time constant and this DOC, I was expecting a JIT crash for this example

def exo2(x):
  n = x.shape[0]             # how n is know at compile time?
  print("x: ",x, type(x))
  jax.debug.print("x = {}",x)
  print("n: ",n, type(n))
  y = jnp.arange(0.,n)    # <-- crash expected here !
  return jnp.sum(x*y)

When no JIT is activated:

x=jax.random.normal(jax.random.PRNGKey(10),shape=(10,))
exo2(x)

leads to

x:  [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608
 -0.12636082 -0.16064268  0.876103   -0.23290116] <class 'jaxlib.xla_extension.ArrayImpl'>
x = [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608
 -0.12636082 -0.16064268  0.876103   -0.23290116]
n:  10 <class 'int'>
Array(-0.16760588, dtype=float32)

And I am curious to see that jit(exo2)(x)does not crash and leads to

x:  Traced<ShapedArray(float32[10])>with<DynamicJaxprTrace(level=1/0)> <class 'jax.interpreters.partial_eval.DynamicJaxprTracer'>
n:  10 <class 'int'>
x = [ 0.4676388  -0.26684377  1.1685165  -0.48255268 -1.4658067   0.408608
 -0.12636082 -0.16064268  0.876103   -0.23290116]

Array(-0.16760588, dtype=float32)

I was (wrongly) expected that the n variable cannot be estimated at compile time ? I probably missed something.

[jakevdp]

Thanks for the question! Properties of arrays like shapes, sizes, and dtypes are always static within JAX, so they can be used to define other static quantities. Does that make sense?

[jecampagne]

Thanks for the doc link. Does it comes from my questioning?

[soraros]

... So everything looks to be working as expected.

@jakevdp Yeah, I read the implementation. Still find the lost of "phase distinction" unsettling though.

This is from the doc:

A useful pattern is to use numpy for operations that should be static (i.e. done at compile-time),
and use jax.numpy for operations that should be traced (i.e. compiled and executed at run-time).

Should we mention using NumPy functions don't provide guarantee (say, raise when not possible at compile time) for staticness here?

[jakevdp]

One strength of that doc is its conciseness; I don't think adding a side-conversation about numpy implementation details would add much value.

[jecampagne]

Oh. yes for sure, I was asking if this doc has been triggered by our discussion? I'm glad if it is the case. Now I feel more comfortable to explain why there was NO crash... Thanks.

[jakevdp]

Hi @jecampagne – my answer was in response to @soraros. Regarding your question, no this doc is a few years old, but it's a frequent question so it's nice to have a good writeup to point to 😀