jax

7569e800 - revert #4277 (google failure) (#4281)

revert #4277 (google failure) (#4281) * revert #4277 (google failure) Some downstream user is relying on the rank of stax's biases being 1. * only revert one change