3 ms·Jax and Equinox – An Introduction1 points by sebg 2y agofithisux 2y agokaf_layer = jax.vmap(KAF(n, key)) is wrong, collab has something else kaf_layer = KAF(n, key=key)