Make JAX funcify more robust to constant inputs

[ricardoV94]

When compiling models for the JAX backend we sometimes see failures (i.e., concretization errors) due to operations like Reshape, even when the shape input is a constant and should be known at compilation time. I haven't found a MWE example that reproduces this issue, and I am not quite sure what causes it just yet, hence why this is a discussion.

This is the current jax_funcify_Reshape:

aesara/aesara/link/jax/dispatch.py

Lines 701 to 706 in 2450186

	@jax_funcify.register(Reshape)
	def jax_funcify_Reshape(op, **kwargs):
	def reshape(x, shape):
	return jnp.reshape(x, shape)
	return reshape

There are two hackish patterns I have seen/used to go around the failures:

1. Check if shape is constant and return a function that completely ignores the shape input

@jax_funcify.register(Reshape)
def jax_funcify_Reshape(op, node, **kwargs):

    shape = node.inputs[1]
    if isinstance(shape, Constant):        
        constant_shape = shape.data
        def reshape(x, _):
            return jnp.reshape(x, constant_shape)

    else:        
        def reshape(x, shape):        
            return jnp.reshape(x, shape)

    return reshape

2. Use jax.ensure_compile_time_eval (this was used by @ferrine recently):

@jax_funcify.register(at.basic_opt.Reshape)
def jax_funcify_Reshape(op, **kwargs):
    # there are corner cases where we need to patch the linker for existing Ops
    def reshape(x, shape):
        with jax.ensure_compile_time_eval():
            x_size = x.size
            s = x_size // jnp.prod(shape, where=shape != -1)
            shape = jnp.where(shape != -1, shape, s)
            return jnp.reshape(x, shape)
    return reshape

In both cases the idea is to tell JAX that the reshape is a safe operation that can be done at compilation time (or at least can be attempted).

I am not very familiar with JAX yet, to know what gotchas we might be getting into here. If someone is interested in exploring this issue, it seems like there is quite some room to increase the coverage of JAX-able graphs we can generate with Aesara by doing some more work inside those jax_funcify* functions.

Related to #43, #68, #182, #684

[brandonwillard]

1. Check if shape is constant and return a function that completely ignores the shape input

@jax_funcify.register(Reshape)
def jax_funcify_Reshape(op, node, **kwargs):

    shape = node.inputs[1]
    if isinstance(shape, Constant):        
        constant_shape = shape.data
        def reshape(x, _):
            return jnp.reshape(x, constant_shape)

    else:        
        def reshape(x, shape):        
            return jnp.reshape(x, shape)

    return reshape

This is something we should consider doing in a more general way (e.g. in aesara.link.basic.JITLinker and/or aesara.link.utils.fgraph_to_python).

One question, and potential complication, that arises when trying to automate this: do we preserve the original Op/node arguments when they're constant or do we elide them? If we take the latter approach, the generated JIT-able functions' inputs won't match the corresponding Ops' design—or the corresponding Apply nodes' inputs. If we don't, then a bunch of unused, constant arguments are passed around. If the target language can remove those arguments itself, then it's not an issue, though.

[brandonwillard]

N.B. Numba/LLVM performs constant propagation sufficiently enough for us to not need manual Constant handling/optimizations in our transpilation code.