flax
2eb7b985 - Add "extracting intermediates" pattern doc (#746)

Commit
5 years ago
Add "extracting intermediates" pattern doc (#746) Adds a page under docs/patterns/ for ways to extract intermediate values from a model. One of the ways suggests using `put_variable`, which this commit also adds back. I propose that this format replace our previous HOWTO system, and we use this to document many of our other patterns, and rewrite all HOWTOs this way. Thus, I am looking for a very careful review. TODO: - [ ] Add unit tests. (The doc itself is an integration test) - [ ] Add a docstring to `put_variable` - [ ] Clone submodules when assigning to `self.layers` in `Sequential` so that we can safely reuse the same `CNN` instance. ADDITIONAL CHANGES ------------------ * Remove the HOWTO system. It was clunky and added maintenance burden to all example editors. Instead, this document contains `testcode` blocks which means that it is tested directly from within the doc, without depending on any examples. * Add support for tests embedded within docs. This allows for writing docs that "feel like" notebooks, but are easier to edit, and allow hiding imports and other initialization in `testsetup` * Allow taking in non-frozen collections in `apply` (resolves #719). QUESTIONS --------- 1. Is it weird that we need to pass in an empty intermediates collection if we want to use `put_variable`? We already specify that we expect changes to that collection via `mutable=['intermediates']`. We could try making `mutable=['intermediates']` add an additional intermediates collection if it's not there but then we'd have to theoretically also support `mutable=[True, 'intermediates']`. Another alternative is to add another option like mutable with an explicit list of collection names that are "output-only" 2. Is the `put_variable` implementation, which allows for either receiving intermediate outputs or not based on whether you pass in an `intermediates` variable collection good? Otherwise, people would have to thread in additional attributes to keep track or whether they want to return intermediate outputs or not. 3. Is the pattern in `Sequential` weird? In `setup`, we run `self.layers = self.submodules` to bind unbound modules. It's hard for users to understand this. I haven't thought too deeply about a possible better solution.
References
Author
Committer
Parents
Loading