Marc van Zee
Marc van Zee
I spend quite some time figuring out why code in a large codebase was so slow, only to find out that `jit` was disabled throughout the entire project. This was...
# What does this PR do? Fixes # (issue) ## Checklist - [ ] This PR fixes a minor issue (e.g.: typo or small bug) or improves the docs (you...
Example code: ``` model = Model(1) #@jax.jit def eval_step(model, params, batch): logits = model.apply({'params': params}, batch['X']) return compute_metrics(logits, batch['y']) def eval_model(model, params, test_ds): metrics = eval_step(model, params, test_ds) metrics =...
Jraph uses `jax.tree_multimap`, which is deprecated and causes our tests to fail. Therefore I added a filterwarning in `pytest.ini`. We should remove this filter warning once Jraph fixes this.
This issue tracks which HOWTOs we would like to add. ## Process - If you like to work on HOWTO, please create a PR for it and mention this issue...
Following #1416 we should further improve Flax Modules on ReadTheDocs. (Content below from @andsteing) We now have an `Parameters` section with only names/types:  And then the attributes after the...
Github Discussions doesn't always seem to work well with Flax, which means some questions that are answerable do not show up in Google search. Here is an example: Google search...
Add the JAX standardize alias for normalize as well, see #2088. This is currently blocked since JAX just added these so we should add them once JAX rolls out a...
Currently our HOWTOs often have nice side-by-side views, but they also hide some code. It would be useful for users to be able to just copy/paste all code from a...
In `setup()`, submodules are expected to "know their name" when they are assigned to `self`, which is implemented by overriding `__setattr__`. This can cause problems when appending modules to a...