Kernel Shape Calculation about neural-tangents HOT 4 CLOSED

Hi @DarrenZhang01, I'm happy to take this one. The idea here is that during computation of the kernel, we would like to know the shape of the intermediate (pre-)activations in the finite version of the computation. To do this, we use JAX's abstract evaluation machinery to infer the shapes using the init_fn without actually instantiating any parameters. Here akey is an abstract version of key that retains only shape and dtype information. If you want to know more, you might want to check out one of the JAX talks (there is a great one by Skye that's recorded somewhere) where they explain how tracing works.

from neural-tangents.

Thanks very much Sam @sschoenholz ! I see that akey is an abstract level (ShapedArray) of key representation. If I understand correctly, when akey serves as the input for abstract_eval_fun in _propagate_shape, it is only used as a PartialVal object in generating the Jaxpr?

from neural-tangents.

I think that's basically correct. As a technical point, I believe JAX doesn't instantiate the jaxpr explicitly, but evaluates the shape while tracing the jaxpr.

from neural-tangents.

I see. Thanks, Sam!

from neural-tangents.