Andreas Steiner

In #1916 we saw some strange updates in the metrics when doing a no-op refactoring. This needs more investigation, maybe also a replacement of the dataset.

Correctly compute metrics over entire evalset/testset

3

When computing metrics for the evaluation/test datasets, the entire set should be processed to get the correct metrics. Our examples often compute metrics with the following pattern: ```python def compute_metrics(logits,...

Publishing pretrained examples on a public GCS bucket would help users to quickly get started with trained models for finetuning, for using them directly for inference etc.

Updating linen_examples/vae

2

This issue tracks updates the `vae` example to follow practices outlined in #231. - [x] Port to linen API - once ported all subsequent changes should be done in `linen_examples/vae`...

Updating examples/nlp_seq

6

This issue tracks updates the `nlp_seq` example to follow practices outlined in #231. - [ ] Port to linen API - once ported all subsequent changes should be done in...

Directly calling `module.setup()` should raise an exception

2

`flax.linen.Module.setup()` should not be called directly because it needs a `flax.linen.Scope` to be set up properly. Since #653 there is no more exception risen when a user inadvertently calls `flax.linen.Module.setup()`...

Adding a new transposed convolution function (similar to `torch.nn.ConvTranspose2d()`)

8

Adding a transposed convolution as proposed in https://github.com/google/jax/pull/5772 would also be very useful when porting models from PyTorch to Flax (as in #1848).

Zeroshot evaluation CIFAR100 in Colab: Results ~0.5% below values reported in paper

3

Hi Thanks so much for providing this repository and the notebooks! I'm debugging diffs in the zeroshot evaluation results from a JAX port of this repository ([`scenic.projects.baselines.clip`](https://github.com/google-research/scenic/tree/main/scenic/projects/baselines/clip)) and as part...

Showcase "transfer learning" with imagenet example

1

The existing [`examples/imagenet/notebook.ipynb`](https://colab.research.google.com/github/google/flax/blob/main/examples/imagenet/imagenet.ipynb) could be extended to showcase how to do transfer learning, i.e. use a pre-trained imagenet model, remove some layers, and then train with some other layers frozen....

Return key from `traverse_util.Traversal.iterate()`

6

Currently `traverse_util.Traversal.iterate()` only returns the traversed values. In some cases we need access to the traversed keys as well, for example when checking that two `ModelParamTraversal` do not overlap in...