Reorganize the jax files to separate the layers.

Also do the padding with xla ops directly in the lowering function. Not sure if this is better really, but it makes the primitive make more sense.

Author: nshepperd