Andreas Steiner

I find myself often copying setup code (e.g. for a simple MNIST-like example) when trying out new features. It would be very handy to have specially crafted Snippets for Flax...

This should fix #88.

When computing the log probabilities for ```python prompts = ( 'the sky is blue', 'the sky is pink', 'the sky is bacon', ) ``` I get very different values, depending...

TPU Colab fails with `AttributeError: module 'jax.tree_util' has no attribute 'register_pytree_with_keys_class'`

1

On a fresh Colab TPU runtime, we have: ``` >>> import flax AttributeError Traceback (most recent call last) [](https://localhost:8080/#) in () 1 # Verify we can import everything. ----> 2...