Avital Oliver

Linen: Consider raising an error when reading variables from submodule before initialization?

10

Within a module that uses shape-inference (as most of the built-in Linen modules do), this code is fine: ``` class MyModule(nn.Module): def __call__(self, x): conv = nn.Conv(features=3) y = conv(x)...

Try a first attempt at a full re-haul of the RTD landing page with more explicit calls to action. This is a starting point for future exploration. See a preview...

cc @levskaya who wrote one that may be ready to release

Ideally we want both type checkers to work correctly at least for the library. Also mypy should offer a speedup _Originally posted by @jheek in https://github.com/google/flax/discussions/685#discussioncomment-140141_

Example tests to run in parallel

8

Right now we run each example's tests sequentially in `tests/run_all_tests.sh` (this is due to https://github.com/pytest-dev/pytest/issues/3151). Because of that, we can't run different example tests in parallel (though pytest can do...

Add "DotGetter" style variable and submodule accessors?

1

This would allow accessing variables defined in a `nn.compact`-wrapped method via `module.kernel` rather than `module.variables['param']['kernel']` Notably, we already have a check that disallows defining variables with the same name in...

Convert Linen design tests into real unit tests

2

See https://github.com/google/flax/tree/master/linen_examples/linen_design_test

Deprecate flax.jax_utils.replicate

1

`flax.jax_utils.replicate` should no longer be necessary and superceded by `jax.device_put_replicated`. We need to update our tests and use-cases accordingly, and add a clear and useful deprecation warning (or, perhaps even...

How to deal with shared submodules?

8

Hi @cgarciae -- treex is cool, in particular I really like the idea of using different classes for "variable kinds" and how it plays pretty nicely with typing (though I...