Pass less state to jax generate function#1398

mattdangerw merged 1 commit intokeras-team:masterkeras-team/keras-hub:masterfrom

mattdangerw:jax-generate-statelessmattdangerw/keras-hub:jax-generate-statelessCopy head branch name to clipboard