Handle the new JAX bind-params API#42

n-gao merged 2 commits intomicrosoft:mainmicrosoft/folx:mainfrom

netket:codex/fix-get-bind-params-upstreamnetket/folx:codex/fix-get-bind-params-upstreamCopy head branch name to clipboard