Added initializers doc #2776

chiamp · 2023-01-07T03:05:22Z

Resolves #2749 and #1386.

Created initializers documentation. View the doc here.

review-notebook-app · 2023-01-07T03:05:26Z

Check out this pull request on

See visual diffs & provide feedback on Jupyter Notebooks.

Powered by ReviewNB

codecov-commenter · 2023-01-07T03:20:01Z

Codecov Report

Merging #2776 (eed5cd9) into main (6bccee3) will not change coverage.
The diff coverage is n/a.

@@           Coverage Diff           @@
##             main    #2776   +/-   ##
=======================================
  Coverage   81.24%   81.24%           
=======================================
  Files          53       53           
  Lines        5663     5663           
=======================================
  Hits         4601     4601           
  Misses       1062     1062

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

marcvanzee

Thanks for adding this guide, really nice! I suppose the next step is to improve docstrings? For instance, linking to this guide from docstrings seems useful.

docs/guides/initializers.md

marcvanzee · 2023-01-09T10:58:08Z

docs/guides/initializers.md

+
+`Initializers` are functions that can be passed as optional arguments to the kernel initializer (`kernel_init`) and the bias initializer (`bias_init`) if you want to specify how the parameters of a Module layer are initialized. A full list of Flax initializers can be found [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html#module-flax.linen.initializers), and are in fact, the same as the [JAX initializers](https://jax.readthedocs.io/en/latest/jax.nn.initializers.html).
+
+The default kernel initializer is [`flax.linen.initializers.lecun_normal`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.lecun_normal.html) and the default bias initializer is [`flax.linen.initializers.zeros`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.zeros.html).


It would be useful if you could explain why we are using these initializers. There was some confusion about this before (see #215). Maybe open a thread on flax-core to ask?

docs/guides/initializers.md

marcvanzee · 2023-01-09T11:36:45Z

docs/guides/initializers.md

+
+++ {"id": "3Kglqd9vuxTG"}
+
+To maintain consistency, all `Initializer` functions that are passed to the `kernel_init` and `bias_init` arguments **must follow the function signature: `[PRNGKey, Shape, Dtype] -> Array`**. Most functions in the [Flax initializer list](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html#module-flax.linen.initializers) are **builder functions** and build an `Initializer` function that follows this function signature. The two exceptions are [`flax.linen.initializers.zeros`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.zeros.html) and [`flax.linen.initializers.ones`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.ones.html), which are already `Initializer` functions that follow the function signature. This is why in the above example, we must call `lecun_normal()` to build an `Initializer` function, whereas we can directly use `zeros` since it's already an `Initializer` function.


Thinking of "prefer simple APIs over more documentation", I am wondering whether it wouldn't be easier to just add two fake builder functions for zeros and ones as follows:

def zeros(dtype: DTypeLikeInexact = jnp.float_) -> Array: def init(key: KeyArray, shape: core.Shape, dtype: DTypeLikeInexact = dtype) -> Array: del key return jnp.zeros(shape, dtype)

Then we don't need this complicated explanation and people can simply use all initializers consistently. WDYT?

I think this is a great idea! It would make things much more consistent. The question is what we should do with the original flax.linen.initializers.zeros and flax.linen.initializers.ones. If we replace them with the fake builder functions you suggested, then wouldn't this break code using these initializers? On the other hand if we leave them, then I think users may get confused on what the difference is; i.e. which they should use, which ones to call versus which ones to use explicitly, etc.

docs/guides/initializers.md

marcvanzee · 2023-01-09T11:41:25Z

docs/guides/initializers.md

+
+++ {"id": "S4X_xHHk-b4V"}
+
+## `Initializer` restrictions for `bias_init`


Nice! Maybe we should link to this from our bias_init docstrings as well?

Do you mean all docstrings that contain a description for bias_init (like here)? Something like: bias_init: initializer function for the bias. To see restrictions on valid initializers, refer to our guide: https://flax.readthedocs.io/en/latest/guides/initializers.html#initializer-restrictions-for-bias-init

chiamp force-pushed the initializers_doc branch from 8564591 to bbbbd36 Compare January 7, 2023 03:13

chiamp requested a review from marcvanzee January 7, 2023 03:22

chiamp self-assigned this Jan 7, 2023

chiamp force-pushed the initializers_doc branch from bbbbd36 to 6c112a0 Compare January 7, 2023 23:12

marcvanzee requested changes Jan 9, 2023

View reviewed changes

chiamp mentioned this pull request Jan 12, 2023

Added builder functions for zeros and ones initializers #2790

Merged

chiamp force-pushed the initializers_doc branch from 6c112a0 to d03db7f Compare January 24, 2023 00:04

Added initializers doc

eed5cd9

chiamp force-pushed the initializers_doc branch from d03db7f to eed5cd9 Compare January 24, 2023 00:07

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Added initializers doc #2776

Added initializers doc #2776

chiamp commented Jan 7, 2023 •

edited

Loading

review-notebook-app bot commented Jan 7, 2023

codecov-commenter commented Jan 7, 2023 •

edited

Loading

marcvanzee left a comment

marcvanzee Jan 9, 2023

marcvanzee Jan 9, 2023

chiamp Jan 10, 2023 •

edited

Loading

marcvanzee Jan 9, 2023

chiamp Jan 10, 2023 •

edited

Loading


		`Initializers` are functions that can be passed as optional arguments to the kernel initializer (`kernel_init`) and the bias initializer (`bias_init`) if you want to specify how the parameters of a Module layer are initialized. A full list of Flax initializers can be found [here](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html#module-flax.linen.initializers), and are in fact, the same as the [JAX initializers](https://jax.readthedocs.io/en/latest/jax.nn.initializers.html).

		The default kernel initializer is [`flax.linen.initializers.lecun_normal`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.lecun_normal.html) and the default bias initializer is [`flax.linen.initializers.zeros`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.zeros.html).


		+++ {"id": "3Kglqd9vuxTG"}

		To maintain consistency, all `Initializer` functions that are passed to the `kernel_init` and `bias_init` arguments must follow the function signature: `[PRNGKey, Shape, Dtype] -> Array`. Most functions in the [Flax initializer list](https://flax.readthedocs.io/en/latest/api_reference/flax.linen.html#module-flax.linen.initializers) are builder functions and build an `Initializer` function that follows this function signature. The two exceptions are [`flax.linen.initializers.zeros`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.zeros.html) and [`flax.linen.initializers.ones`](https://flax.readthedocs.io/en/latest/api_reference/_autosummary/flax.linen.initializers.ones.html), which are already `Initializer` functions that follow the function signature. This is why in the above example, we must call `lecun_normal()` to build an `Initializer` function, whereas we can directly use `zeros` since it's already an `Initializer` function.


		+++ {"id": "S4X_xHHk-b4V"}

		## `Initializer` restrictions for `bias_init`

Added initializers doc #2776

Are you sure you want to change the base?

Added initializers doc #2776

Conversation

chiamp commented Jan 7, 2023 • edited Loading

review-notebook-app bot commented Jan 7, 2023

codecov-commenter commented Jan 7, 2023 • edited Loading

Codecov Report

marcvanzee left a comment

Choose a reason for hiding this comment

marcvanzee Jan 9, 2023

Choose a reason for hiding this comment

marcvanzee Jan 9, 2023

Choose a reason for hiding this comment

chiamp Jan 10, 2023 • edited Loading

Choose a reason for hiding this comment

marcvanzee Jan 9, 2023

Choose a reason for hiding this comment

chiamp Jan 10, 2023 • edited Loading

Choose a reason for hiding this comment

chiamp commented Jan 7, 2023 •

edited

Loading

codecov-commenter commented Jan 7, 2023 •

edited

Loading

chiamp Jan 10, 2023 •

edited

Loading

chiamp Jan 10, 2023 •

edited

Loading