AOT path doesn't typecheck DCE'd arguments

[froystig]

>>> import jax
>>> import jax.numpy as jnp
>>> x, y1, y2 = jnp.ones(4), jnp.ones(4), jnp.ones(1)

>>> f = lambda x, y: x + 1. if y.shape[0] > 2 else x + 2.
>>> f(x, y1)
Array([2., 2., 2., 2.], dtype=float32)
>>> f(x, y2)
Array([3., 3., 3., 3.], dtype=float32)

>>> g = jax.jit(f).lower(x, y1).compile()
>>> g(x, y1)
Array([2., 2., 2., 2.], dtype=float32)
>>> g(x, y2)
Array([2., 2., 2., 2.], dtype=float32)

The second call to g should raise an error.

[froystig]

Shardings are also not checked on DCE's operands. If we ever exposed shardings from tracers, we could exhibit a similar issue with mismatched shardings in place of shapes (avals) above.

Ahead of that, we might still want to cautiously error because, on the JIT path, different shardings result in re-compilation, and part of the purpose of the AOT API is to guard against re-compilation. In other words, we could adapt an example like the one above that triggers a side effect (such as a print statement) differently between JIT and AOT.

[froystig]

Finally, the same goes for layouts once we've made them available.