We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
When jax_enable_x64 is set Adam promotes everything to float64

### Problem you have encountered:

When `jax_enable_x64` is set, Adam's `apply_gradient` method will promote all float32 arrays to float64, potentially unexpectedly degrading performance.

This is due to jax's wonky type promotion semantics. The offending line is:
https://github.com/google/flax/blob/3e36db3e5e3b8e6e1777d612f270e7948238aa9c/flax/optim/adam.py#L82

which promotes like:
```python
jnp.array([0], dtype=jnp.int32) + 1. # == DeviceArray([1.], dtype=float64)
```
and then cascades from there promoting everything to float64

### What you expected to happen:

Arrays should retain their dtypes on optimizer updates.

### Logs, error messages, etc:


### Steps to reproduce:

```python
from jax.config import config
config.update("jax_enable_x64", True)

import jax.numpy as jnp
import flax

opt = flax.optim.Adam(1e-3).create(
        {"x": jnp.zeros(10, dtype=jnp.float32)})

assert opt.target["x"].dtype == jnp.float32

opt = opt.apply_gradient({"x": jnp.zeros(10, dtype=jnp.float32)})

# This fails, since dtype was promoted to float64
assert opt.target["x"].dtype == jnp.float32
```

--- 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/optim/lamb.py`
Content:
```
1 # Copyright 2021 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 from .. import struct
16 
17 from jax import lax
18 import jax.numpy as jnp
19 
20 import numpy as onp
21 
22 from .base import OptimizerDef
23 
24 @struct.dataclass
25 class _LAMBHyperParams:
26   learning_rate: onp.ndarray
27   beta1: onp.ndarray
28   beta2: onp.ndarray
29   weight_decay: onp.ndarray
30   eps: onp.ndarray
31 
32 
33 @struct.dataclass
34 class _LAMBParamState:
35   grad_ema: onp.ndarray
36   grad_sq_ema: onp.ndarray
37 
38 
39 class LAMB(OptimizerDef):
40   """Layerwise adaptive moments for batch (LAMB) optimizer.
41 
42   See https://arxiv.org/abs/1904.00962
43   """
44 
45   def __init__(self, learning_rate=None, beta1=0.9, beta2=0.999, weight_decay=0,
46                eps=1e-6):
47     """Constructor for the LAMB optimizer.
48 
49     Args:
50       learning_rate: the step size used to update the parameters.
51       beta1: the coefficient used for the moving average of the gradient
52         (default: 0.9).
53       beta2: the coefficient used for the moving average of the squared gradient
54         (default: 0.999).
55       weight_decay: weight decay coefficient to apply
56       eps: epsilon used for Adam update computation (default: 1e-6).
57     """
58 
59     hyper_params = _LAMBHyperParams(
60         learning_rate, beta1, beta2, weight_decay, eps)
61     super().__init__(hyper_params)
62 
63   def init_param_state(self, param):
64     return _LAMBParamState(jnp.zeros_like(param), jnp.zeros_like(param))
65 
66   def apply_param_gradient(self, step, hyper_params, param, state, grad):
67     assert hyper_params.learning_rate is not None, 'no learning rate provided.'
68     beta1 = hyper_params.beta1
69     beta2 = hyper_params.beta2
70     weight_decay = hyper_params.weight_decay
71     learning_rate = hyper_params.learning_rate
72 
73     grad_sq = lax.square(grad)
74     grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad
75     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq
76 
77     t = step + 1.
78     grad_ema_corr = grad_ema / (1. - beta1 ** t)
79     grad_sq_ema_corr = grad_sq_ema / (1. - beta2 ** t)
80 
81     update = grad_ema_corr / (jnp.sqrt(grad_sq_ema_corr) + hyper_params.eps)
82 
83     if weight_decay != 0.0:
84       update += weight_decay * param
85 
86     param_norm = jnp.linalg.norm(param)
87     update_norm = jnp.linalg.norm(update)
88     trust_ratio = jnp.where(
89         param_norm + update_norm > 0., param_norm / update_norm, 1.)
90 
91     new_param = param - trust_ratio * learning_rate * update
92     new_state = _LAMBParamState(grad_ema, grad_sq_ema)
93     return new_param, new_state
94
```

Path: `flax/optim/adam.py`
Content:
```
1 # Copyright 2021 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 from .. import struct
16 
17 import jax.numpy as jnp
18 from jax import lax
19 
20 import numpy as onp
21 
22 from .base import OptimizerDef
23 
24 
25 @struct.dataclass
26 class _AdamHyperParams:
27   learning_rate: onp.ndarray
28   beta1: onp.ndarray
29   beta2: onp.ndarray
30   eps: onp.ndarray
31   weight_decay: onp.ndarray
32 
33 
34 @struct.dataclass
35 class _AdamParamState:
36   grad_ema: onp.ndarray
37   grad_sq_ema: onp.ndarray
38 
39 
40 class Adam(OptimizerDef):
41   """Adam optimizer.
42 
43   Implements Adam - a stochastic gradient descent method (SGD) that computes
44   individual adaptive learning rates for different parameters from estimates of
45   first- and second-order moments of the gradients. 
46   
47   Reference: [Adam: A Method
48   for Stochastic Optimization](https://arxiv.org/abs/1412.6980v8) (Kingma and
49   Ba, 2014).
50 
51   Attributes:
52     learning_rate: The learning rate — the step size used to update the
53       parameters.
54     beta1: The exponentian decay rate for the 1st moment estimates. The
55       coefficient used to calculate the first moments of the gradients (the
56       moving average of the gradient) (default: 0.9).
57     beta2: The exponentian decay rate for the 2nd moment estimates. The
58       coefficient used to calculate the second moments of the gradients (the
59       moving average of the gradient magnitude) (default: 0.999).
60     eps: A small scalar added to the gradient magnitude estimate to improve
61       numerical stability (default: 1e-8).
62     weight_decay: The learning rate decay (default: 0.0).
63   """
64 
65   def __init__(self,
66                learning_rate=None,
67                beta1=0.9,
68                beta2=0.999,
69                eps=1e-8,
70                weight_decay=0.0):
71     """Constructor for the Adam optimizer.
72 
73     Args:
74       learning_rate: The step size used to update the parameters.
75       beta1: The coefficient used for the moving average of the
76         gradient (default: 0.9).
77       beta2: The coefficient used for the moving average of the
78         gradient magnitude (default: 0.999).
79       eps: The term added to the gradient magnitude estimate for
80         numerical stability (default: 1e-8).
81       weight_decay: AdamW style weight decay rate
82         (relative to learning rate) (default: 0.0).
83     """
84     hyper_params = _AdamHyperParams(learning_rate, beta1, beta2, eps,
85                                     weight_decay)
86     super().__init__(hyper_params)
87 
88   def init_param_state(self, param):
89     return _AdamParamState(jnp.zeros_like(param), jnp.zeros_like(param))
90 
91   def apply_param_gradient(self, step, hyper_params, param, state, grad):
92     assert hyper_params.learning_rate is not None, 'no learning rate provided.'
93     beta1 = hyper_params.beta1
94     beta2 = hyper_params.beta2
95     weight_decay = hyper_params.weight_decay
96     grad_sq = lax.square(grad)
97     grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad
98     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq
99 
100     # bias correction
101     t = step + 1.
102     grad_ema_corr = grad_ema / (1 - beta1 ** t)
103     grad_sq_ema_corr = grad_sq_ema / (1 - beta2 ** t)
104 
105     denom = jnp.sqrt(grad_sq_ema_corr) + hyper_params.eps
106     new_param = param - hyper_params.learning_rate * grad_ema_corr / denom
107     new_param -= hyper_params.learning_rate * weight_decay * param
108     new_state = _AdamParamState(grad_ema, grad_sq_ema)
109     return new_param, new_state
110
```
--- 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/optim/adam.py b/flax/optim/adam.py
--- a/flax/optim/adam.py
+++ b/flax/optim/adam.py
@@ -98,7 +98,7 @@
     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq
 
     # bias correction
-    t = step + 1.
+    t = jnp.array(step + 1, lax.dtype(param.dtype))
     grad_ema_corr = grad_ema / (1 - beta1 ** t)
     grad_sq_ema_corr = grad_sq_ema / (1 - beta2 ** t)
 
diff --git a/flax/optim/lamb.py b/flax/optim/lamb.py
--- a/flax/optim/lamb.py
+++ b/flax/optim/lamb.py
@@ -74,7 +74,7 @@
     grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad
     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq
 
-    t = step + 1.
+    t = jnp.array(step + 1, lax.dtype(param.dtype))
     grad_ema_corr = grad_ema / (1. - beta1 ** t)
     grad_sq_ema_corr = grad_sq_ema / (1. - beta2 ** t)

{"golden_diff": "diff --git a/flax/optim/adam.py b/flax/optim/adam.py\n--- a/flax/optim/adam.py\n+++ b/flax/optim/adam.py\n@@ -98,7 +98,7 @@\n     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq\n \n     # bias correction\n-    t = step + 1.\n+    t = jnp.array(step + 1, lax.dtype(param.dtype))\n     grad_ema_corr = grad_ema / (1 - beta1 ** t)\n     grad_sq_ema_corr = grad_sq_ema / (1 - beta2 ** t)\n \ndiff --git a/flax/optim/lamb.py b/flax/optim/lamb.py\n--- a/flax/optim/lamb.py\n+++ b/flax/optim/lamb.py\n@@ -74,7 +74,7 @@\n     grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad\n     grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq\n \n-    t = step + 1.\n+    t = jnp.array(step + 1, lax.dtype(param.dtype))\n     grad_ema_corr = grad_ema / (1. - beta1 ** t)\n     grad_sq_ema_corr = grad_sq_ema / (1. - beta2 ** t)\n", "issue": "When jax_enable_x64 is set Adam promotes everything to float64\n\r\n### Problem you have encountered:\r\n\r\nWhen `jax_enable_x64` is set, Adam's `apply_gradient` method will promote all float32 arrays to float64, potentially unexpectedly degrading performance.\r\n\r\nThis is due to jax's wonky type promotion semantics. The offending line is:\r\nhttps://github.com/google/flax/blob/3e36db3e5e3b8e6e1777d612f270e7948238aa9c/flax/optim/adam.py#L82\r\n\r\nwhich promotes like:\r\n```python\r\njnp.array([0], dtype=jnp.int32) + 1. # == DeviceArray([1.], dtype=float64)\r\n```\r\nand then cascades from there promoting everything to float64\r\n\r\n### What you expected to happen:\r\n\r\nArrays should retain their dtypes on optimizer updates.\r\n\r\n### Logs, error messages, etc:\r\n\r\n\r\n### Steps to reproduce:\r\n\r\n```python\r\nfrom jax.config import config\r\nconfig.update(\"jax_enable_x64\", True)\r\n\r\nimport jax.numpy as jnp\r\nimport flax\r\n\r\nopt = flax.optim.Adam(1e-3).create(\r\n        {\"x\": jnp.zeros(10, dtype=jnp.float32)})\r\n\r\nassert opt.target[\"x\"].dtype == jnp.float32\r\n\r\nopt = opt.apply_gradient({\"x\": jnp.zeros(10, dtype=jnp.float32)})\r\n\r\n# This fails, since dtype was promoted to float64\r\nassert opt.target[\"x\"].dtype == jnp.float32\r\n```\n", "before_files": [{"content": "# Copyright 2021 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\nfrom .. import struct\n\nfrom jax import lax\nimport jax.numpy as jnp\n\nimport numpy as onp\n\nfrom .base import OptimizerDef\n\n@struct.dataclass\nclass _LAMBHyperParams:\n  learning_rate: onp.ndarray\n  beta1: onp.ndarray\n  beta2: onp.ndarray\n  weight_decay: onp.ndarray\n  eps: onp.ndarray\n\n\n@struct.dataclass\nclass _LAMBParamState:\n  grad_ema: onp.ndarray\n  grad_sq_ema: onp.ndarray\n\n\nclass LAMB(OptimizerDef):\n  \"\"\"Layerwise adaptive moments for batch (LAMB) optimizer.\n\n  See https://arxiv.org/abs/1904.00962\n  \"\"\"\n\n  def __init__(self, learning_rate=None, beta1=0.9, beta2=0.999, weight_decay=0,\n               eps=1e-6):\n    \"\"\"Constructor for the LAMB optimizer.\n\n    Args:\n      learning_rate: the step size used to update the parameters.\n      beta1: the coefficient used for the moving average of the gradient\n        (default: 0.9).\n      beta2: the coefficient used for the moving average of the squared gradient\n        (default: 0.999).\n      weight_decay: weight decay coefficient to apply\n      eps: epsilon used for Adam update computation (default: 1e-6).\n    \"\"\"\n\n    hyper_params = _LAMBHyperParams(\n        learning_rate, beta1, beta2, weight_decay, eps)\n    super().__init__(hyper_params)\n\n  def init_param_state(self, param):\n    return _LAMBParamState(jnp.zeros_like(param), jnp.zeros_like(param))\n\n  def apply_param_gradient(self, step, hyper_params, param, state, grad):\n    assert hyper_params.learning_rate is not None, 'no learning rate provided.'\n    beta1 = hyper_params.beta1\n    beta2 = hyper_params.beta2\n    weight_decay = hyper_params.weight_decay\n    learning_rate = hyper_params.learning_rate\n\n    grad_sq = lax.square(grad)\n    grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad\n    grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq\n\n    t = step + 1.\n    grad_ema_corr = grad_ema / (1. - beta1 ** t)\n    grad_sq_ema_corr = grad_sq_ema / (1. - beta2 ** t)\n\n    update = grad_ema_corr / (jnp.sqrt(grad_sq_ema_corr) + hyper_params.eps)\n\n    if weight_decay != 0.0:\n      update += weight_decay * param\n\n    param_norm = jnp.linalg.norm(param)\n    update_norm = jnp.linalg.norm(update)\n    trust_ratio = jnp.where(\n        param_norm + update_norm > 0., param_norm / update_norm, 1.)\n\n    new_param = param - trust_ratio * learning_rate * update\n    new_state = _LAMBParamState(grad_ema, grad_sq_ema)\n    return new_param, new_state\n", "path": "flax/optim/lamb.py"}, {"content": "# Copyright 2021 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\nfrom .. import struct\n\nimport jax.numpy as jnp\nfrom jax import lax\n\nimport numpy as onp\n\nfrom .base import OptimizerDef\n\n\n@struct.dataclass\nclass _AdamHyperParams:\n  learning_rate: onp.ndarray\n  beta1: onp.ndarray\n  beta2: onp.ndarray\n  eps: onp.ndarray\n  weight_decay: onp.ndarray\n\n\n@struct.dataclass\nclass _AdamParamState:\n  grad_ema: onp.ndarray\n  grad_sq_ema: onp.ndarray\n\n\nclass Adam(OptimizerDef):\n  \"\"\"Adam optimizer.\n\n  Implements Adam - a stochastic gradient descent method (SGD) that computes\n  individual adaptive learning rates for different parameters from estimates of\n  first- and second-order moments of the gradients. \n  \n  Reference: [Adam: A Method\n  for Stochastic Optimization](https://arxiv.org/abs/1412.6980v8) (Kingma and\n  Ba, 2014).\n\n  Attributes:\n    learning_rate: The learning rate \u2014 the step size used to update the\n      parameters.\n    beta1: The exponentian decay rate for the 1st moment estimates. The\n      coefficient used to calculate the first moments of the gradients (the\n      moving average of the gradient) (default: 0.9).\n    beta2: The exponentian decay rate for the 2nd moment estimates. The\n      coefficient used to calculate the second moments of the gradients (the\n      moving average of the gradient magnitude) (default: 0.999).\n    eps: A small scalar added to the gradient magnitude estimate to improve\n      numerical stability (default: 1e-8).\n    weight_decay: The learning rate decay (default: 0.0).\n  \"\"\"\n\n  def __init__(self,\n               learning_rate=None,\n               beta1=0.9,\n               beta2=0.999,\n               eps=1e-8,\n               weight_decay=0.0):\n    \"\"\"Constructor for the Adam optimizer.\n\n    Args:\n      learning_rate: The step size used to update the parameters.\n      beta1: The coefficient used for the moving average of the\n        gradient (default: 0.9).\n      beta2: The coefficient used for the moving average of the\n        gradient magnitude (default: 0.999).\n      eps: The term added to the gradient magnitude estimate for\n        numerical stability (default: 1e-8).\n      weight_decay: AdamW style weight decay rate\n        (relative to learning rate) (default: 0.0).\n    \"\"\"\n    hyper_params = _AdamHyperParams(learning_rate, beta1, beta2, eps,\n                                    weight_decay)\n    super().__init__(hyper_params)\n\n  def init_param_state(self, param):\n    return _AdamParamState(jnp.zeros_like(param), jnp.zeros_like(param))\n\n  def apply_param_gradient(self, step, hyper_params, param, state, grad):\n    assert hyper_params.learning_rate is not None, 'no learning rate provided.'\n    beta1 = hyper_params.beta1\n    beta2 = hyper_params.beta2\n    weight_decay = hyper_params.weight_decay\n    grad_sq = lax.square(grad)\n    grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad\n    grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq\n\n    # bias correction\n    t = step + 1.\n    grad_ema_corr = grad_ema / (1 - beta1 ** t)\n    grad_sq_ema_corr = grad_sq_ema / (1 - beta2 ** t)\n\n    denom = jnp.sqrt(grad_sq_ema_corr) + hyper_params.eps\n    new_param = param - hyper_params.learning_rate * grad_ema_corr / denom\n    new_param -= hyper_params.learning_rate * weight_decay * param\n    new_state = _AdamParamState(grad_ema, grad_sq_ema)\n    return new_param, new_state\n", "path": "flax/optim/adam.py"}], "after_files": [{"content": "# Copyright 2021 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\nfrom .. import struct\n\nfrom jax import lax\nimport jax.numpy as jnp\n\nimport numpy as onp\n\nfrom .base import OptimizerDef\n\n@struct.dataclass\nclass _LAMBHyperParams:\n  learning_rate: onp.ndarray\n  beta1: onp.ndarray\n  beta2: onp.ndarray\n  weight_decay: onp.ndarray\n  eps: onp.ndarray\n\n\n@struct.dataclass\nclass _LAMBParamState:\n  grad_ema: onp.ndarray\n  grad_sq_ema: onp.ndarray\n\n\nclass LAMB(OptimizerDef):\n  \"\"\"Layerwise adaptive moments for batch (LAMB) optimizer.\n\n  See https://arxiv.org/abs/1904.00962\n  \"\"\"\n\n  def __init__(self, learning_rate=None, beta1=0.9, beta2=0.999, weight_decay=0,\n               eps=1e-6):\n    \"\"\"Constructor for the LAMB optimizer.\n\n    Args:\n      learning_rate: the step size used to update the parameters.\n      beta1: the coefficient used for the moving average of the gradient\n        (default: 0.9).\n      beta2: the coefficient used for the moving average of the squared gradient\n        (default: 0.999).\n      weight_decay: weight decay coefficient to apply\n      eps: epsilon used for Adam update computation (default: 1e-6).\n    \"\"\"\n\n    hyper_params = _LAMBHyperParams(\n        learning_rate, beta1, beta2, weight_decay, eps)\n    super().__init__(hyper_params)\n\n  def init_param_state(self, param):\n    return _LAMBParamState(jnp.zeros_like(param), jnp.zeros_like(param))\n\n  def apply_param_gradient(self, step, hyper_params, param, state, grad):\n    assert hyper_params.learning_rate is not None, 'no learning rate provided.'\n    beta1 = hyper_params.beta1\n    beta2 = hyper_params.beta2\n    weight_decay = hyper_params.weight_decay\n    learning_rate = hyper_params.learning_rate\n\n    grad_sq = lax.square(grad)\n    grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad\n    grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq\n\n    t = jnp.array(step + 1, lax.dtype(param.dtype))\n    grad_ema_corr = grad_ema / (1. - beta1 ** t)\n    grad_sq_ema_corr = grad_sq_ema / (1. - beta2 ** t)\n\n    update = grad_ema_corr / (jnp.sqrt(grad_sq_ema_corr) + hyper_params.eps)\n\n    if weight_decay != 0.0:\n      update += weight_decay * param\n\n    param_norm = jnp.linalg.norm(param)\n    update_norm = jnp.linalg.norm(update)\n    trust_ratio = jnp.where(\n        param_norm + update_norm > 0., param_norm / update_norm, 1.)\n\n    new_param = param - trust_ratio * learning_rate * update\n    new_state = _LAMBParamState(grad_ema, grad_sq_ema)\n    return new_param, new_state\n", "path": "flax/optim/lamb.py"}, {"content": "# Copyright 2021 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\nfrom .. import struct\n\nimport jax.numpy as jnp\nfrom jax import lax\n\nimport numpy as onp\n\nfrom .base import OptimizerDef\n\n\n@struct.dataclass\nclass _AdamHyperParams:\n  learning_rate: onp.ndarray\n  beta1: onp.ndarray\n  beta2: onp.ndarray\n  eps: onp.ndarray\n  weight_decay: onp.ndarray\n\n\n@struct.dataclass\nclass _AdamParamState:\n  grad_ema: onp.ndarray\n  grad_sq_ema: onp.ndarray\n\n\nclass Adam(OptimizerDef):\n  \"\"\"Adam optimizer.\n\n  Implements Adam - a stochastic gradient descent method (SGD) that computes\n  individual adaptive learning rates for different parameters from estimates of\n  first- and second-order moments of the gradients. \n  \n  Reference: [Adam: A Method\n  for Stochastic Optimization](https://arxiv.org/abs/1412.6980v8) (Kingma and\n  Ba, 2014).\n\n  Attributes:\n    learning_rate: The learning rate \u2014 the step size used to update the\n      parameters.\n    beta1: The exponentian decay rate for the 1st moment estimates. The\n      coefficient used to calculate the first moments of the gradients (the\n      moving average of the gradient) (default: 0.9).\n    beta2: The exponentian decay rate for the 2nd moment estimates. The\n      coefficient used to calculate the second moments of the gradients (the\n      moving average of the gradient magnitude) (default: 0.999).\n    eps: A small scalar added to the gradient magnitude estimate to improve\n      numerical stability (default: 1e-8).\n    weight_decay: The learning rate decay (default: 0.0).\n  \"\"\"\n\n  def __init__(self,\n               learning_rate=None,\n               beta1=0.9,\n               beta2=0.999,\n               eps=1e-8,\n               weight_decay=0.0):\n    \"\"\"Constructor for the Adam optimizer.\n\n    Args:\n      learning_rate: The step size used to update the parameters.\n      beta1: The coefficient used for the moving average of the\n        gradient (default: 0.9).\n      beta2: The coefficient used for the moving average of the\n        gradient magnitude (default: 0.999).\n      eps: The term added to the gradient magnitude estimate for\n        numerical stability (default: 1e-8).\n      weight_decay: AdamW style weight decay rate\n        (relative to learning rate) (default: 0.0).\n    \"\"\"\n    hyper_params = _AdamHyperParams(learning_rate, beta1, beta2, eps,\n                                    weight_decay)\n    super().__init__(hyper_params)\n\n  def init_param_state(self, param):\n    return _AdamParamState(jnp.zeros_like(param), jnp.zeros_like(param))\n\n  def apply_param_gradient(self, step, hyper_params, param, state, grad):\n    assert hyper_params.learning_rate is not None, 'no learning rate provided.'\n    beta1 = hyper_params.beta1\n    beta2 = hyper_params.beta2\n    weight_decay = hyper_params.weight_decay\n    grad_sq = lax.square(grad)\n    grad_ema = beta1 * state.grad_ema + (1. - beta1) * grad\n    grad_sq_ema = beta2 * state.grad_sq_ema + (1. - beta2) * grad_sq\n\n    # bias correction\n    t = jnp.array(step + 1, lax.dtype(param.dtype))\n    grad_ema_corr = grad_ema / (1 - beta1 ** t)\n    grad_sq_ema_corr = grad_sq_ema / (1 - beta2 ** t)\n\n    denom = jnp.sqrt(grad_sq_ema_corr) + hyper_params.eps\n    new_param = param - hyper_params.learning_rate * grad_ema_corr / denom\n    new_param -= hyper_params.learning_rate * weight_decay * param\n    new_state = _AdamParamState(grad_ema, grad_sq_ema)\n    return new_param, new_state\n", "path": "flax/optim/adam.py"}]}