jax

8ffb9417 - fix stax initialization rng bug, remove temp file