We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[Accessibility] Enable EPUB output on ReadTheDocs
Currently there is only HTML output [enabled](https://readthedocs.org/projects/flax/downloads/). It would be great if EPUB and PDF could also be enabled. 

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flax/struct.py`
Content:
```
1 # Copyright 2022 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Utilities for defining custom classes that can be used with jax transformations.
16 """
17 
18 import dataclasses
19 import typing
20 from typing import TypeVar, Callable, Tuple, Union, Any
21 
22 from . import serialization
23 
24 import jax
25 from typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet
26 
27 
28 _T = TypeVar("_T")
29 
30 
31 def field(pytree_node=True, **kwargs):
32   return dataclasses.field(metadata={'pytree_node': pytree_node}, **kwargs)
33 
34 
35 @dataclass_transform(field_descriptors=(field,))
36 def dataclass(clz: _T) -> _T:
37   """Create a class which can be passed to functional transformations.
38 
39   NOTE: Inherit from ``PyTreeNode`` instead to avoid type checking issues when
40   using PyType.
41 
42   Jax transformations such as `jax.jit` and `jax.grad` require objects that are
43   immutable and can be mapped over using the `jax.tree_util` methods.
44   The `dataclass` decorator makes it easy to define custom classes that can be
45   passed safely to Jax. For example::
46 
47     from flax import struct
48 
49     @struct.dataclass
50     class Model:
51       params: Any
52       # use pytree_node=False to indicate an attribute should not be touched
53       # by Jax transformations.
54       apply_fn: FunctionType = struct.field(pytree_node=False)
55 
56       def __apply__(self, *args):
57         return self.apply_fn(*args)
58 
59     model = Model(params, apply_fn)
60 
61     model.params = params_b  # Model is immutable. This will raise an error.
62     model_b = model.replace(params=params_b)  # Use the replace method instead.
63 
64     # This class can now be used safely in Jax to compute gradients w.r.t. the
65     # parameters.
66     model = Model(params, apply_fn)
67     model_grad = jax.grad(some_loss_fn)(model)
68 
69   Note that dataclasses have an auto-generated ``__init__`` where
70   the arguments of the constructor and the attributes of the created
71   instance match 1:1. This correspondence is what makes these objects
72   valid containers that work with JAX transformations and
73   more generally the `jax.tree_util` library.
74 
75   Sometimes a "smart constructor" is desired, for example because
76   some of the attributes can be (optionally) derived from others.
77   The way to do this with Flax dataclasses is to make a static or
78   class method that provides the smart constructor.
79   This way the simple constructor used by `jax.tree_util` is
80   preserved. Consider the following example::
81 
82     @struct.dataclass
83     class DirectionAndScaleKernel:
84       direction: Array
85       scale: Array
86 
87       @classmethod
88       def create(cls, kernel):
89         scale = jax.numpy.linalg.norm(kernel, axis=0, keepdims=True)
90         directin = direction / scale
91         return cls(direction, scale)
92 
93   Args:
94     clz: the class that will be transformed by the decorator.
95   Returns:
96     The new class.
97   """
98   # check if already a flax dataclass
99   if '_flax_dataclass' in clz.__dict__:
100     return clz
101       
102   data_clz = dataclasses.dataclass(frozen=True)(clz)
103   meta_fields = []
104   data_fields = []
105   for field_info in dataclasses.fields(data_clz):
106     is_pytree_node = field_info.metadata.get('pytree_node', True)
107     if is_pytree_node:
108       data_fields.append(field_info.name)
109     else:
110       meta_fields.append(field_info.name)
111 
112   def replace(self, **updates):
113     """"Returns a new object replacing the specified fields with new values."""
114     return dataclasses.replace(self, **updates)
115 
116   data_clz.replace = replace
117 
118   def iterate_clz(x):
119     meta = tuple(getattr(x, name) for name in meta_fields)
120     data = tuple(getattr(x, name) for name in data_fields)
121     return data, meta
122 
123   def clz_from_iterable(meta, data):
124     meta_args = tuple(zip(meta_fields, meta))
125     data_args = tuple(zip(data_fields, data))
126     kwargs = dict(meta_args + data_args)
127     return data_clz(**kwargs)
128 
129   jax.tree_util.register_pytree_node(data_clz,
130                                      iterate_clz,
131                                      clz_from_iterable)
132 
133   if tuple(map(int, jax.version.__version__.split('.'))) >= (0, 3, 1):
134     def keypaths(_):
135       return [jax.tree_util.AttributeKeyPathEntry(name) for name in data_fields]
136     jax.tree_util.register_keypaths(data_clz, keypaths)
137 
138   def to_state_dict(x):
139     state_dict = {name: serialization.to_state_dict(getattr(x, name))
140                   for name in data_fields}
141     return state_dict
142 
143   def from_state_dict(x, state):
144     """Restore the state of a data class."""
145     state = state.copy()  # copy the state so we can pop the restored fields.
146     updates = {}
147     for name in data_fields:
148       if name not in state:
149         raise ValueError(f'Missing field {name} in state dict while restoring'
150                          f' an instance of {clz.__name__}')
151       value = getattr(x, name)
152       value_state = state.pop(name)
153       updates[name] = serialization.from_state_dict(value, value_state)
154     if state:
155       names = ','.join(state.keys())
156       raise ValueError(f'Unknown field(s) "{names}" in state dict while'
157                        f' restoring an instance of {clz.__name__}')
158     return x.replace(**updates)
159 
160   serialization.register_serialization_state(
161       data_clz, to_state_dict, from_state_dict)
162 
163   # add a _flax_dataclass flag to distinguish from regular dataclasses
164   data_clz._flax_dataclass = True
165 
166   return data_clz
167 
168 
169 TNode = TypeVar('TNode', bound='PyTreeNode')
170 
171 
172 @dataclass_transform(field_descriptors=(field,))
173 class PyTreeNode:
174   """Base class for dataclasses that should act like a JAX pytree node.
175 
176   See ``flax.struct.dataclass`` for the ``jax.tree_util`` behavior.
177   This base class additionally avoids type checking errors when using PyType.
178 
179   Example::
180 
181     from flax import struct
182 
183     class Model(struct.PyTreeNode):
184       params: Any
185       # use pytree_node=False to indicate an attribute should not be touched
186       # by Jax transformations.
187       apply_fn: FunctionType = struct.field(pytree_node=False)
188 
189       def __apply__(self, *args):
190         return self.apply_fn(*args)
191 
192     model = Model(params, apply_fn)
193 
194     model.params = params_b  # Model is immutable. This will raise an error.
195     model_b = model.replace(params=params_b)  # Use the replace method instead.
196 
197     # This class can now be used safely in Jax to compute gradients w.r.t. the
198     # parameters.
199     model = Model(params, apply_fn)
200     model_grad = jax.grad(some_loss_fn)(model)
201 
202   """
203 
204   def __init_subclass__(cls):
205     dataclass(cls)  # pytype: disable=wrong-arg-types
206 
207   def __init__(self, *args, **kwargs):
208     # stub for pytype
209     raise NotImplementedError
210 
211   def replace(self: TNode, **overrides) -> TNode:
212     # stub for pytype
213     raise NotImplementedError
214
```

Path: `flax/linen/stochastic.py`
Content:
```
1 # Copyright 2022 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Stochastic modules."""
16 
17 from typing import Optional, Sequence
18 
19 from flax.linen.module import compact
20 from flax.linen.module import merge_param
21 from flax.linen.module import Module
22 from jax import lax
23 from jax import random
24 import jax.numpy as jnp
25 
26 
27 class Dropout(Module):
28   """Create a dropout layer.
29 
30     Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
31     to include an RNG seed named `'dropout'`. For example::
32     
33       model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
34 
35     Attributes:
36       rate: the dropout probability.  (_not_ the keep rate!)
37       broadcast_dims: dimensions that will share the same dropout mask
38       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
39         masked, whereas if true, no mask is applied and the inputs are returned
40         as is.
41   """
42   rate: float
43   broadcast_dims: Sequence[int] = ()
44   deterministic: Optional[bool] = None
45 
46   @compact
47   def __call__(self, inputs, deterministic: Optional[bool] = None):
48     """Applies a random dropout mask to the input.
49 
50     Args:
51       inputs: the inputs that should be randomly masked.
52       deterministic: if false the inputs are scaled by `1 / (1 - rate)` and
53         masked, whereas if true, no mask is applied and the inputs are returned
54         as is.
55 
56     Returns:
57       The masked inputs reweighted to preserve mean.
58     """
59     deterministic = merge_param(
60         'deterministic', self.deterministic, deterministic)
61     if self.rate == 0.:
62       return inputs
63     # Prevent gradient NaNs in 1.0 edge-case.
64     if self.rate == 1.0:
65       return jnp.zeros_like(inputs)
66     keep_prob = 1. - self.rate
67     if deterministic:
68       return inputs
69     else:
70       rng = self.make_rng('dropout')
71       broadcast_shape = list(inputs.shape)
72       for dim in self.broadcast_dims:
73         broadcast_shape[dim] = 1
74       mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)
75       mask = jnp.broadcast_to(mask, inputs.shape)
76       return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))
77
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py
--- a/flax/linen/stochastic.py
+++ b/flax/linen/stochastic.py
@@ -29,7 +29,7 @@
 
     Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure
     to include an RNG seed named `'dropout'`. For example::
-    
+
       model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`
 
     Attributes:
diff --git a/flax/struct.py b/flax/struct.py
--- a/flax/struct.py
+++ b/flax/struct.py
@@ -98,7 +98,7 @@
   # check if already a flax dataclass
   if '_flax_dataclass' in clz.__dict__:
     return clz
-      
+
   data_clz = dataclasses.dataclass(frozen=True)(clz)
   meta_fields = []
   data_fields = []

{"golden_diff": "diff --git a/flax/linen/stochastic.py b/flax/linen/stochastic.py\n--- a/flax/linen/stochastic.py\n+++ b/flax/linen/stochastic.py\n@@ -29,7 +29,7 @@\n \n     Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure\n     to include an RNG seed named `'dropout'`. For example::\n-    \n+\n       model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`\n \n     Attributes:\ndiff --git a/flax/struct.py b/flax/struct.py\n--- a/flax/struct.py\n+++ b/flax/struct.py\n@@ -98,7 +98,7 @@\n   # check if already a flax dataclass\n   if '_flax_dataclass' in clz.__dict__:\n     return clz\n-      \n+\n   data_clz = dataclasses.dataclass(frozen=True)(clz)\n   meta_fields = []\n   data_fields = []\n", "issue": "[Accessibility] Enable EPUB output on ReadTheDocs\nCurrently there is only HTML output [enabled](https://readthedocs.org/projects/flax/downloads/). It would be great if EPUB and PDF could also be enabled. \n", "before_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for defining custom classes that can be used with jax transformations.\n\"\"\"\n\nimport dataclasses\nimport typing\nfrom typing import TypeVar, Callable, Tuple, Union, Any\n\nfrom . import serialization\n\nimport jax\nfrom typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet\n\n\n_T = TypeVar(\"_T\")\n\n\ndef field(pytree_node=True, **kwargs):\n  return dataclasses.field(metadata={'pytree_node': pytree_node}, **kwargs)\n\n\n@dataclass_transform(field_descriptors=(field,))\ndef dataclass(clz: _T) -> _T:\n  \"\"\"Create a class which can be passed to functional transformations.\n\n  NOTE: Inherit from ``PyTreeNode`` instead to avoid type checking issues when\n  using PyType.\n\n  Jax transformations such as `jax.jit` and `jax.grad` require objects that are\n  immutable and can be mapped over using the `jax.tree_util` methods.\n  The `dataclass` decorator makes it easy to define custom classes that can be\n  passed safely to Jax. For example::\n\n    from flax import struct\n\n    @struct.dataclass\n    class Model:\n      params: Any\n      # use pytree_node=False to indicate an attribute should not be touched\n      # by Jax transformations.\n      apply_fn: FunctionType = struct.field(pytree_node=False)\n\n      def __apply__(self, *args):\n        return self.apply_fn(*args)\n\n    model = Model(params, apply_fn)\n\n    model.params = params_b  # Model is immutable. This will raise an error.\n    model_b = model.replace(params=params_b)  # Use the replace method instead.\n\n    # This class can now be used safely in Jax to compute gradients w.r.t. the\n    # parameters.\n    model = Model(params, apply_fn)\n    model_grad = jax.grad(some_loss_fn)(model)\n\n  Note that dataclasses have an auto-generated ``__init__`` where\n  the arguments of the constructor and the attributes of the created\n  instance match 1:1. This correspondence is what makes these objects\n  valid containers that work with JAX transformations and\n  more generally the `jax.tree_util` library.\n\n  Sometimes a \"smart constructor\" is desired, for example because\n  some of the attributes can be (optionally) derived from others.\n  The way to do this with Flax dataclasses is to make a static or\n  class method that provides the smart constructor.\n  This way the simple constructor used by `jax.tree_util` is\n  preserved. Consider the following example::\n\n    @struct.dataclass\n    class DirectionAndScaleKernel:\n      direction: Array\n      scale: Array\n\n      @classmethod\n      def create(cls, kernel):\n        scale = jax.numpy.linalg.norm(kernel, axis=0, keepdims=True)\n        directin = direction / scale\n        return cls(direction, scale)\n\n  Args:\n    clz: the class that will be transformed by the decorator.\n  Returns:\n    The new class.\n  \"\"\"\n  # check if already a flax dataclass\n  if '_flax_dataclass' in clz.__dict__:\n    return clz\n      \n  data_clz = dataclasses.dataclass(frozen=True)(clz)\n  meta_fields = []\n  data_fields = []\n  for field_info in dataclasses.fields(data_clz):\n    is_pytree_node = field_info.metadata.get('pytree_node', True)\n    if is_pytree_node:\n      data_fields.append(field_info.name)\n    else:\n      meta_fields.append(field_info.name)\n\n  def replace(self, **updates):\n    \"\"\"\"Returns a new object replacing the specified fields with new values.\"\"\"\n    return dataclasses.replace(self, **updates)\n\n  data_clz.replace = replace\n\n  def iterate_clz(x):\n    meta = tuple(getattr(x, name) for name in meta_fields)\n    data = tuple(getattr(x, name) for name in data_fields)\n    return data, meta\n\n  def clz_from_iterable(meta, data):\n    meta_args = tuple(zip(meta_fields, meta))\n    data_args = tuple(zip(data_fields, data))\n    kwargs = dict(meta_args + data_args)\n    return data_clz(**kwargs)\n\n  jax.tree_util.register_pytree_node(data_clz,\n                                     iterate_clz,\n                                     clz_from_iterable)\n\n  if tuple(map(int, jax.version.__version__.split('.'))) >= (0, 3, 1):\n    def keypaths(_):\n      return [jax.tree_util.AttributeKeyPathEntry(name) for name in data_fields]\n    jax.tree_util.register_keypaths(data_clz, keypaths)\n\n  def to_state_dict(x):\n    state_dict = {name: serialization.to_state_dict(getattr(x, name))\n                  for name in data_fields}\n    return state_dict\n\n  def from_state_dict(x, state):\n    \"\"\"Restore the state of a data class.\"\"\"\n    state = state.copy()  # copy the state so we can pop the restored fields.\n    updates = {}\n    for name in data_fields:\n      if name not in state:\n        raise ValueError(f'Missing field {name} in state dict while restoring'\n                         f' an instance of {clz.__name__}')\n      value = getattr(x, name)\n      value_state = state.pop(name)\n      updates[name] = serialization.from_state_dict(value, value_state)\n    if state:\n      names = ','.join(state.keys())\n      raise ValueError(f'Unknown field(s) \"{names}\" in state dict while'\n                       f' restoring an instance of {clz.__name__}')\n    return x.replace(**updates)\n\n  serialization.register_serialization_state(\n      data_clz, to_state_dict, from_state_dict)\n\n  # add a _flax_dataclass flag to distinguish from regular dataclasses\n  data_clz._flax_dataclass = True\n\n  return data_clz\n\n\nTNode = TypeVar('TNode', bound='PyTreeNode')\n\n\n@dataclass_transform(field_descriptors=(field,))\nclass PyTreeNode:\n  \"\"\"Base class for dataclasses that should act like a JAX pytree node.\n\n  See ``flax.struct.dataclass`` for the ``jax.tree_util`` behavior.\n  This base class additionally avoids type checking errors when using PyType.\n\n  Example::\n\n    from flax import struct\n\n    class Model(struct.PyTreeNode):\n      params: Any\n      # use pytree_node=False to indicate an attribute should not be touched\n      # by Jax transformations.\n      apply_fn: FunctionType = struct.field(pytree_node=False)\n\n      def __apply__(self, *args):\n        return self.apply_fn(*args)\n\n    model = Model(params, apply_fn)\n\n    model.params = params_b  # Model is immutable. This will raise an error.\n    model_b = model.replace(params=params_b)  # Use the replace method instead.\n\n    # This class can now be used safely in Jax to compute gradients w.r.t. the\n    # parameters.\n    model = Model(params, apply_fn)\n    model_grad = jax.grad(some_loss_fn)(model)\n\n  \"\"\"\n\n  def __init_subclass__(cls):\n    dataclass(cls)  # pytype: disable=wrong-arg-types\n\n  def __init__(self, *args, **kwargs):\n    # stub for pytype\n    raise NotImplementedError\n\n  def replace(self: TNode, **overrides) -> TNode:\n    # stub for pytype\n    raise NotImplementedError\n", "path": "flax/struct.py"}, {"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Stochastic modules.\"\"\"\n\nfrom typing import Optional, Sequence\n\nfrom flax.linen.module import compact\nfrom flax.linen.module import merge_param\nfrom flax.linen.module import Module\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\nclass Dropout(Module):\n  \"\"\"Create a dropout layer.\n\n    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure\n    to include an RNG seed named `'dropout'`. For example::\n    \n      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`\n\n    Attributes:\n      rate: the dropout probability.  (_not_ the keep rate!)\n      broadcast_dims: dimensions that will share the same dropout mask\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n  \"\"\"\n  rate: float\n  broadcast_dims: Sequence[int] = ()\n  deterministic: Optional[bool] = None\n\n  @compact\n  def __call__(self, inputs, deterministic: Optional[bool] = None):\n    \"\"\"Applies a random dropout mask to the input.\n\n    Args:\n      inputs: the inputs that should be randomly masked.\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n\n    Returns:\n      The masked inputs reweighted to preserve mean.\n    \"\"\"\n    deterministic = merge_param(\n        'deterministic', self.deterministic, deterministic)\n    if self.rate == 0.:\n      return inputs\n    # Prevent gradient NaNs in 1.0 edge-case.\n    if self.rate == 1.0:\n      return jnp.zeros_like(inputs)\n    keep_prob = 1. - self.rate\n    if deterministic:\n      return inputs\n    else:\n      rng = self.make_rng('dropout')\n      broadcast_shape = list(inputs.shape)\n      for dim in self.broadcast_dims:\n        broadcast_shape[dim] = 1\n      mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)\n      mask = jnp.broadcast_to(mask, inputs.shape)\n      return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n", "path": "flax/linen/stochastic.py"}], "after_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utilities for defining custom classes that can be used with jax transformations.\n\"\"\"\n\nimport dataclasses\nimport typing\nfrom typing import TypeVar, Callable, Tuple, Union, Any\n\nfrom . import serialization\n\nimport jax\nfrom typing_extensions import dataclass_transform  # pytype: disable=not-supported-yet\n\n\n_T = TypeVar(\"_T\")\n\n\ndef field(pytree_node=True, **kwargs):\n  return dataclasses.field(metadata={'pytree_node': pytree_node}, **kwargs)\n\n\n@dataclass_transform(field_descriptors=(field,))\ndef dataclass(clz: _T) -> _T:\n  \"\"\"Create a class which can be passed to functional transformations.\n\n  NOTE: Inherit from ``PyTreeNode`` instead to avoid type checking issues when\n  using PyType.\n\n  Jax transformations such as `jax.jit` and `jax.grad` require objects that are\n  immutable and can be mapped over using the `jax.tree_util` methods.\n  The `dataclass` decorator makes it easy to define custom classes that can be\n  passed safely to Jax. For example::\n\n    from flax import struct\n\n    @struct.dataclass\n    class Model:\n      params: Any\n      # use pytree_node=False to indicate an attribute should not be touched\n      # by Jax transformations.\n      apply_fn: FunctionType = struct.field(pytree_node=False)\n\n      def __apply__(self, *args):\n        return self.apply_fn(*args)\n\n    model = Model(params, apply_fn)\n\n    model.params = params_b  # Model is immutable. This will raise an error.\n    model_b = model.replace(params=params_b)  # Use the replace method instead.\n\n    # This class can now be used safely in Jax to compute gradients w.r.t. the\n    # parameters.\n    model = Model(params, apply_fn)\n    model_grad = jax.grad(some_loss_fn)(model)\n\n  Note that dataclasses have an auto-generated ``__init__`` where\n  the arguments of the constructor and the attributes of the created\n  instance match 1:1. This correspondence is what makes these objects\n  valid containers that work with JAX transformations and\n  more generally the `jax.tree_util` library.\n\n  Sometimes a \"smart constructor\" is desired, for example because\n  some of the attributes can be (optionally) derived from others.\n  The way to do this with Flax dataclasses is to make a static or\n  class method that provides the smart constructor.\n  This way the simple constructor used by `jax.tree_util` is\n  preserved. Consider the following example::\n\n    @struct.dataclass\n    class DirectionAndScaleKernel:\n      direction: Array\n      scale: Array\n\n      @classmethod\n      def create(cls, kernel):\n        scale = jax.numpy.linalg.norm(kernel, axis=0, keepdims=True)\n        directin = direction / scale\n        return cls(direction, scale)\n\n  Args:\n    clz: the class that will be transformed by the decorator.\n  Returns:\n    The new class.\n  \"\"\"\n  # check if already a flax dataclass\n  if '_flax_dataclass' in clz.__dict__:\n    return clz\n\n  data_clz = dataclasses.dataclass(frozen=True)(clz)\n  meta_fields = []\n  data_fields = []\n  for field_info in dataclasses.fields(data_clz):\n    is_pytree_node = field_info.metadata.get('pytree_node', True)\n    if is_pytree_node:\n      data_fields.append(field_info.name)\n    else:\n      meta_fields.append(field_info.name)\n\n  def replace(self, **updates):\n    \"\"\"\"Returns a new object replacing the specified fields with new values.\"\"\"\n    return dataclasses.replace(self, **updates)\n\n  data_clz.replace = replace\n\n  def iterate_clz(x):\n    meta = tuple(getattr(x, name) for name in meta_fields)\n    data = tuple(getattr(x, name) for name in data_fields)\n    return data, meta\n\n  def clz_from_iterable(meta, data):\n    meta_args = tuple(zip(meta_fields, meta))\n    data_args = tuple(zip(data_fields, data))\n    kwargs = dict(meta_args + data_args)\n    return data_clz(**kwargs)\n\n  jax.tree_util.register_pytree_node(data_clz,\n                                     iterate_clz,\n                                     clz_from_iterable)\n\n  if tuple(map(int, jax.version.__version__.split('.'))) >= (0, 3, 1):\n    def keypaths(_):\n      return [jax.tree_util.AttributeKeyPathEntry(name) for name in data_fields]\n    jax.tree_util.register_keypaths(data_clz, keypaths)\n\n  def to_state_dict(x):\n    state_dict = {name: serialization.to_state_dict(getattr(x, name))\n                  for name in data_fields}\n    return state_dict\n\n  def from_state_dict(x, state):\n    \"\"\"Restore the state of a data class.\"\"\"\n    state = state.copy()  # copy the state so we can pop the restored fields.\n    updates = {}\n    for name in data_fields:\n      if name not in state:\n        raise ValueError(f'Missing field {name} in state dict while restoring'\n                         f' an instance of {clz.__name__}')\n      value = getattr(x, name)\n      value_state = state.pop(name)\n      updates[name] = serialization.from_state_dict(value, value_state)\n    if state:\n      names = ','.join(state.keys())\n      raise ValueError(f'Unknown field(s) \"{names}\" in state dict while'\n                       f' restoring an instance of {clz.__name__}')\n    return x.replace(**updates)\n\n  serialization.register_serialization_state(\n      data_clz, to_state_dict, from_state_dict)\n\n  # add a _flax_dataclass flag to distinguish from regular dataclasses\n  data_clz._flax_dataclass = True\n\n  return data_clz\n\n\nTNode = TypeVar('TNode', bound='PyTreeNode')\n\n\n@dataclass_transform(field_descriptors=(field,))\nclass PyTreeNode:\n  \"\"\"Base class for dataclasses that should act like a JAX pytree node.\n\n  See ``flax.struct.dataclass`` for the ``jax.tree_util`` behavior.\n  This base class additionally avoids type checking errors when using PyType.\n\n  Example::\n\n    from flax import struct\n\n    class Model(struct.PyTreeNode):\n      params: Any\n      # use pytree_node=False to indicate an attribute should not be touched\n      # by Jax transformations.\n      apply_fn: FunctionType = struct.field(pytree_node=False)\n\n      def __apply__(self, *args):\n        return self.apply_fn(*args)\n\n    model = Model(params, apply_fn)\n\n    model.params = params_b  # Model is immutable. This will raise an error.\n    model_b = model.replace(params=params_b)  # Use the replace method instead.\n\n    # This class can now be used safely in Jax to compute gradients w.r.t. the\n    # parameters.\n    model = Model(params, apply_fn)\n    model_grad = jax.grad(some_loss_fn)(model)\n\n  \"\"\"\n\n  def __init_subclass__(cls):\n    dataclass(cls)  # pytype: disable=wrong-arg-types\n\n  def __init__(self, *args, **kwargs):\n    # stub for pytype\n    raise NotImplementedError\n\n  def replace(self: TNode, **overrides) -> TNode:\n    # stub for pytype\n    raise NotImplementedError\n", "path": "flax/struct.py"}, {"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Stochastic modules.\"\"\"\n\nfrom typing import Optional, Sequence\n\nfrom flax.linen.module import compact\nfrom flax.linen.module import merge_param\nfrom flax.linen.module import Module\nfrom jax import lax\nfrom jax import random\nimport jax.numpy as jnp\n\n\nclass Dropout(Module):\n  \"\"\"Create a dropout layer.\n\n    Note: When using :meth:`Module.apply() <flax.linen.Module.apply>`, make sure\n    to include an RNG seed named `'dropout'`. For example::\n\n      model.apply({'params': params}, inputs=inputs, train=True, rngs={'dropout': dropout_rng})`\n\n    Attributes:\n      rate: the dropout probability.  (_not_ the keep rate!)\n      broadcast_dims: dimensions that will share the same dropout mask\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n  \"\"\"\n  rate: float\n  broadcast_dims: Sequence[int] = ()\n  deterministic: Optional[bool] = None\n\n  @compact\n  def __call__(self, inputs, deterministic: Optional[bool] = None):\n    \"\"\"Applies a random dropout mask to the input.\n\n    Args:\n      inputs: the inputs that should be randomly masked.\n      deterministic: if false the inputs are scaled by `1 / (1 - rate)` and\n        masked, whereas if true, no mask is applied and the inputs are returned\n        as is.\n\n    Returns:\n      The masked inputs reweighted to preserve mean.\n    \"\"\"\n    deterministic = merge_param(\n        'deterministic', self.deterministic, deterministic)\n    if self.rate == 0.:\n      return inputs\n    # Prevent gradient NaNs in 1.0 edge-case.\n    if self.rate == 1.0:\n      return jnp.zeros_like(inputs)\n    keep_prob = 1. - self.rate\n    if deterministic:\n      return inputs\n    else:\n      rng = self.make_rng('dropout')\n      broadcast_shape = list(inputs.shape)\n      for dim in self.broadcast_dims:\n        broadcast_shape[dim] = 1\n      mask = random.bernoulli(rng, p=keep_prob, shape=broadcast_shape)\n      mask = jnp.broadcast_to(mask, inputs.shape)\n      return lax.select(mask, inputs / keep_prob, jnp.zeros_like(inputs))\n", "path": "flax/linen/stochastic.py"}]}