We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
VAE example uses deprecated `optimizer.optimize()`


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `examples/vae/main.py`
Content:
```
1 # Copyright 2020 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 absl import app
16 from absl import flags
17 
18 import jax.numpy as jnp
19 import numpy as np
20 
21 import jax
22 from jax import random
23 
24 from flax import nn
25 from flax import optim
26 
27 import tensorflow_datasets as tfds
28 
29 from utils import save_image
30 
31 
32 FLAGS = flags.FLAGS
33 
34 flags.DEFINE_float(
35   'learning_rate', default=1e-3,
36   help=('The leanring rate for the Adam optimizer')
37 )
38 
39 flags.DEFINE_integer(
40   'batch_size', default=128,
41   help=('Batch size for training')
42 )
43 
44 flags.DEFINE_integer(
45   'num_epochs', default=10,
46   help=('Number of training epochs')
47 )
48 
49 
50 class Encoder(nn.Module):
51   def apply(self, x):
52     x = nn.Dense(x, 400, name='enc_fc1')
53     x = nn.relu(x)
54     mean_x = nn.Dense(x, 20, name='enc_fc21')
55     logvar_x = nn.Dense(x, 20, name='enc_fc22')
56     return mean_x, logvar_x
57 
58 
59 class Decoder(nn.Module):
60   def apply(self, z):
61     z = nn.Dense(z, 400, name='dec_fc1')
62     z = nn.relu(z)
63     z = nn.Dense(z, 784, name='dec_fc2')
64     return z
65 
66 
67 class VAE(nn.Module):
68   def apply(self, x):
69     mean, logvar = Encoder(x, name='encoder')
70     z = reparameterize(mean, logvar)
71     recon_x = self._created_decoder()(z)
72     return recon_x, mean, logvar
73 
74   @nn.module_method
75   def generate(self, z):
76     params = self.get_param('decoder')
77     return nn.sigmoid(Decoder.call(params, z))
78 
79   @nn.module_method
80   def generate_one_liner(self, z):
81     return nn.sigmoid(Decoder(z, name='decoder'))
82 
83   @nn.module_method
84   def generate_shared(self, z):
85     return nn.sigmoid(self._created_decoder()(z))
86 
87   def _created_decoder(self):
88     return Decoder.shared(name='decoder')
89 
90 
91 def reparameterize(mean, logvar):
92   std = jnp.exp(0.5 * logvar)
93   eps = np.random.normal(size=logvar.shape)
94   return mean + eps * std
95 
96 
97 @jax.vmap
98 def kl_divergence(mean, logvar):
99   return - 0.5 * jnp.sum(1 + logvar - jnp.power(mean, 2) - jnp.exp(logvar))
100 
101 
102 @jax.vmap
103 def binary_cross_entropy_with_logits(logits, labels):
104   return - jnp.sum(labels * nn.log_sigmoid(logits) + (1 - labels) * (nn.log_sigmoid(logits) - logits))
105 
106 
107 def compute_metrics(recon_x, x, mean, logvar):
108   bce_loss = binary_cross_entropy_with_logits(recon_x, x)
109   kld_loss = kl_divergence(mean, logvar)
110   return {'bce': jnp.mean(bce_loss), 'kld': jnp.mean(kld_loss), 'loss': jnp.mean(bce_loss + kld_loss)}
111 
112 
113 @jax.jit
114 def train_step(optimizer, batch):
115   def loss_fn(model):
116     x = batch['image']
117     recon_x, mean, logvar = model(x)
118 
119     bce_loss = binary_cross_entropy_with_logits(recon_x, x)
120     kld_loss = kl_divergence(mean, logvar)
121     loss = jnp.mean(bce_loss + kld_loss)
122     return loss, recon_x
123   optimizer, _, _ = optimizer.optimize(loss_fn)
124   return optimizer
125 
126 
127 @jax.jit
128 def eval(model, eval_ds, z):
129   xs = eval_ds['image'] / 255.0
130   xs = xs.reshape(-1, 784)
131   recon_xs, mean, logvar = model(xs)
132 
133   comparison = jnp.concatenate([xs[:8].reshape(-1, 28, 28, 1),
134                                 recon_xs[:8].reshape(-1, 28, 28, 1)])
135 
136   generate_xs = model.generate(z)
137   generate_xs = generate_xs.reshape(-1, 28, 28, 1)
138 
139   return compute_metrics(recon_xs, xs, mean, logvar), comparison, generate_xs
140 
141 
142 def main(argv):
143   key = random.PRNGKey(0)
144   train_ds = tfds.load('mnist', split=tfds.Split.TRAIN)
145   train_ds = train_ds.cache().shuffle(1000).batch(FLAGS.batch_size)
146   test_ds = tfds.as_numpy(tfds.load('mnist', split=tfds.Split.TEST, batch_size=-1))
147 
148   _, params = VAE.init_by_shape(key, [((1, 784), jnp.float32)])
149   vae = nn.Model(VAE, params)
150 
151   optimizer = optim.Adam(learning_rate=FLAGS.learning_rate).create(vae)
152 
153   for epoch in range(FLAGS.num_epochs):
154     for batch in tfds.as_numpy(train_ds):
155       batch['image'] = batch['image'].reshape(-1, 784) / 255.0
156       optimizer = train_step(optimizer, batch)
157 
158     z = np.random.normal(size=(64, 20))
159     metrics, comparison, sample = eval(optimizer.target, test_ds, z)
160     save_image(comparison, 'results/reconstruction_' + str(epoch) + '.png', nrow=8)
161     save_image(sample, 'results/sample_' + str(epoch) + '.png', nrow=8)
162 
163     print("eval epoch: {}, loss: {:.4f}, BCE: {:.4f}, KLD: {:.4f}".format(
164       epoch + 1, metrics['loss'], metrics['bce'], metrics['kld']
165     ))
166 
167 
168 if __name__ == '__main__':
169     app.run(main)
170
```
--- 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/examples/vae/main.py b/examples/vae/main.py
--- a/examples/vae/main.py
+++ b/examples/vae/main.py
@@ -120,7 +120,9 @@
     kld_loss = kl_divergence(mean, logvar)
     loss = jnp.mean(bce_loss + kld_loss)
     return loss, recon_x
-  optimizer, _, _ = optimizer.optimize(loss_fn)
+  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)
+  _, grad = grad_fn(optimizer.target)
+  optimizer = optimizer.apply_gradient(grad)
   return optimizer

{"golden_diff": "diff --git a/examples/vae/main.py b/examples/vae/main.py\n--- a/examples/vae/main.py\n+++ b/examples/vae/main.py\n@@ -120,7 +120,9 @@\n     kld_loss = kl_divergence(mean, logvar)\n     loss = jnp.mean(bce_loss + kld_loss)\n     return loss, recon_x\n-  optimizer, _, _ = optimizer.optimize(loss_fn)\n+  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n+  _, grad = grad_fn(optimizer.target)\n+  optimizer = optimizer.apply_gradient(grad)\n   return optimizer\n", "issue": "VAE example uses deprecated `optimizer.optimize()`\n\n", "before_files": [{"content": "# Copyright 2020 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 absl import app\nfrom absl import flags\n\nimport jax.numpy as jnp\nimport numpy as np\n\nimport jax\nfrom jax import random\n\nfrom flax import nn\nfrom flax import optim\n\nimport tensorflow_datasets as tfds\n\nfrom utils import save_image\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_float(\n  'learning_rate', default=1e-3,\n  help=('The leanring rate for the Adam optimizer')\n)\n\nflags.DEFINE_integer(\n  'batch_size', default=128,\n  help=('Batch size for training')\n)\n\nflags.DEFINE_integer(\n  'num_epochs', default=10,\n  help=('Number of training epochs')\n)\n\n\nclass Encoder(nn.Module):\n  def apply(self, x):\n    x = nn.Dense(x, 400, name='enc_fc1')\n    x = nn.relu(x)\n    mean_x = nn.Dense(x, 20, name='enc_fc21')\n    logvar_x = nn.Dense(x, 20, name='enc_fc22')\n    return mean_x, logvar_x\n\n\nclass Decoder(nn.Module):\n  def apply(self, z):\n    z = nn.Dense(z, 400, name='dec_fc1')\n    z = nn.relu(z)\n    z = nn.Dense(z, 784, name='dec_fc2')\n    return z\n\n\nclass VAE(nn.Module):\n  def apply(self, x):\n    mean, logvar = Encoder(x, name='encoder')\n    z = reparameterize(mean, logvar)\n    recon_x = self._created_decoder()(z)\n    return recon_x, mean, logvar\n\n  @nn.module_method\n  def generate(self, z):\n    params = self.get_param('decoder')\n    return nn.sigmoid(Decoder.call(params, z))\n\n  @nn.module_method\n  def generate_one_liner(self, z):\n    return nn.sigmoid(Decoder(z, name='decoder'))\n\n  @nn.module_method\n  def generate_shared(self, z):\n    return nn.sigmoid(self._created_decoder()(z))\n\n  def _created_decoder(self):\n    return Decoder.shared(name='decoder')\n\n\ndef reparameterize(mean, logvar):\n  std = jnp.exp(0.5 * logvar)\n  eps = np.random.normal(size=logvar.shape)\n  return mean + eps * std\n\n\n@jax.vmap\ndef kl_divergence(mean, logvar):\n  return - 0.5 * jnp.sum(1 + logvar - jnp.power(mean, 2) - jnp.exp(logvar))\n\n\n@jax.vmap\ndef binary_cross_entropy_with_logits(logits, labels):\n  return - jnp.sum(labels * nn.log_sigmoid(logits) + (1 - labels) * (nn.log_sigmoid(logits) - logits))\n\n\ndef compute_metrics(recon_x, x, mean, logvar):\n  bce_loss = binary_cross_entropy_with_logits(recon_x, x)\n  kld_loss = kl_divergence(mean, logvar)\n  return {'bce': jnp.mean(bce_loss), 'kld': jnp.mean(kld_loss), 'loss': jnp.mean(bce_loss + kld_loss)}\n\n\n@jax.jit\ndef train_step(optimizer, batch):\n  def loss_fn(model):\n    x = batch['image']\n    recon_x, mean, logvar = model(x)\n\n    bce_loss = binary_cross_entropy_with_logits(recon_x, x)\n    kld_loss = kl_divergence(mean, logvar)\n    loss = jnp.mean(bce_loss + kld_loss)\n    return loss, recon_x\n  optimizer, _, _ = optimizer.optimize(loss_fn)\n  return optimizer\n\n\n@jax.jit\ndef eval(model, eval_ds, z):\n  xs = eval_ds['image'] / 255.0\n  xs = xs.reshape(-1, 784)\n  recon_xs, mean, logvar = model(xs)\n\n  comparison = jnp.concatenate([xs[:8].reshape(-1, 28, 28, 1),\n                                recon_xs[:8].reshape(-1, 28, 28, 1)])\n\n  generate_xs = model.generate(z)\n  generate_xs = generate_xs.reshape(-1, 28, 28, 1)\n\n  return compute_metrics(recon_xs, xs, mean, logvar), comparison, generate_xs\n\n\ndef main(argv):\n  key = random.PRNGKey(0)\n  train_ds = tfds.load('mnist', split=tfds.Split.TRAIN)\n  train_ds = train_ds.cache().shuffle(1000).batch(FLAGS.batch_size)\n  test_ds = tfds.as_numpy(tfds.load('mnist', split=tfds.Split.TEST, batch_size=-1))\n\n  _, params = VAE.init_by_shape(key, [((1, 784), jnp.float32)])\n  vae = nn.Model(VAE, params)\n\n  optimizer = optim.Adam(learning_rate=FLAGS.learning_rate).create(vae)\n\n  for epoch in range(FLAGS.num_epochs):\n    for batch in tfds.as_numpy(train_ds):\n      batch['image'] = batch['image'].reshape(-1, 784) / 255.0\n      optimizer = train_step(optimizer, batch)\n\n    z = np.random.normal(size=(64, 20))\n    metrics, comparison, sample = eval(optimizer.target, test_ds, z)\n    save_image(comparison, 'results/reconstruction_' + str(epoch) + '.png', nrow=8)\n    save_image(sample, 'results/sample_' + str(epoch) + '.png', nrow=8)\n\n    print(\"eval epoch: {}, loss: {:.4f}, BCE: {:.4f}, KLD: {:.4f}\".format(\n      epoch + 1, metrics['loss'], metrics['bce'], metrics['kld']\n    ))\n\n\nif __name__ == '__main__':\n    app.run(main)\n", "path": "examples/vae/main.py"}], "after_files": [{"content": "# Copyright 2020 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 absl import app\nfrom absl import flags\n\nimport jax.numpy as jnp\nimport numpy as np\n\nimport jax\nfrom jax import random\n\nfrom flax import nn\nfrom flax import optim\n\nimport tensorflow_datasets as tfds\n\nfrom utils import save_image\n\n\nFLAGS = flags.FLAGS\n\nflags.DEFINE_float(\n  'learning_rate', default=1e-3,\n  help=('The leanring rate for the Adam optimizer')\n)\n\nflags.DEFINE_integer(\n  'batch_size', default=128,\n  help=('Batch size for training')\n)\n\nflags.DEFINE_integer(\n  'num_epochs', default=10,\n  help=('Number of training epochs')\n)\n\n\nclass Encoder(nn.Module):\n  def apply(self, x):\n    x = nn.Dense(x, 400, name='enc_fc1')\n    x = nn.relu(x)\n    mean_x = nn.Dense(x, 20, name='enc_fc21')\n    logvar_x = nn.Dense(x, 20, name='enc_fc22')\n    return mean_x, logvar_x\n\n\nclass Decoder(nn.Module):\n  def apply(self, z):\n    z = nn.Dense(z, 400, name='dec_fc1')\n    z = nn.relu(z)\n    z = nn.Dense(z, 784, name='dec_fc2')\n    return z\n\n\nclass VAE(nn.Module):\n  def apply(self, x):\n    mean, logvar = Encoder(x, name='encoder')\n    z = reparameterize(mean, logvar)\n    recon_x = self._created_decoder()(z)\n    return recon_x, mean, logvar\n\n  @nn.module_method\n  def generate(self, z):\n    params = self.get_param('decoder')\n    return nn.sigmoid(Decoder.call(params, z))\n\n  @nn.module_method\n  def generate_one_liner(self, z):\n    return nn.sigmoid(Decoder(z, name='decoder'))\n\n  @nn.module_method\n  def generate_shared(self, z):\n    return nn.sigmoid(self._created_decoder()(z))\n\n  def _created_decoder(self):\n    return Decoder.shared(name='decoder')\n\n\ndef reparameterize(mean, logvar):\n  std = jnp.exp(0.5 * logvar)\n  eps = np.random.normal(size=logvar.shape)\n  return mean + eps * std\n\n\n@jax.vmap\ndef kl_divergence(mean, logvar):\n  return - 0.5 * jnp.sum(1 + logvar - jnp.power(mean, 2) - jnp.exp(logvar))\n\n\n@jax.vmap\ndef binary_cross_entropy_with_logits(logits, labels):\n  return - jnp.sum(labels * nn.log_sigmoid(logits) + (1 - labels) * (nn.log_sigmoid(logits) - logits))\n\n\ndef compute_metrics(recon_x, x, mean, logvar):\n  bce_loss = binary_cross_entropy_with_logits(recon_x, x)\n  kld_loss = kl_divergence(mean, logvar)\n  return {'bce': jnp.mean(bce_loss), 'kld': jnp.mean(kld_loss), 'loss': jnp.mean(bce_loss + kld_loss)}\n\n\n@jax.jit\ndef train_step(optimizer, batch):\n  def loss_fn(model):\n    x = batch['image']\n    recon_x, mean, logvar = model(x)\n\n    bce_loss = binary_cross_entropy_with_logits(recon_x, x)\n    kld_loss = kl_divergence(mean, logvar)\n    loss = jnp.mean(bce_loss + kld_loss)\n    return loss, recon_x\n  grad_fn = jax.value_and_grad(loss_fn, has_aux=True)\n  _, grad = grad_fn(optimizer.target)\n  optimizer = optimizer.apply_gradient(grad)\n  return optimizer\n\n\n@jax.jit\ndef eval(model, eval_ds, z):\n  xs = eval_ds['image'] / 255.0\n  xs = xs.reshape(-1, 784)\n  recon_xs, mean, logvar = model(xs)\n\n  comparison = jnp.concatenate([xs[:8].reshape(-1, 28, 28, 1),\n                                recon_xs[:8].reshape(-1, 28, 28, 1)])\n\n  generate_xs = model.generate(z)\n  generate_xs = generate_xs.reshape(-1, 28, 28, 1)\n\n  return compute_metrics(recon_xs, xs, mean, logvar), comparison, generate_xs\n\n\ndef main(argv):\n  key = random.PRNGKey(0)\n  train_ds = tfds.load('mnist', split=tfds.Split.TRAIN)\n  train_ds = train_ds.cache().shuffle(1000).batch(FLAGS.batch_size)\n  test_ds = tfds.as_numpy(tfds.load('mnist', split=tfds.Split.TEST, batch_size=-1))\n\n  _, params = VAE.init_by_shape(key, [((1, 784), jnp.float32)])\n  vae = nn.Model(VAE, params)\n\n  optimizer = optim.Adam(learning_rate=FLAGS.learning_rate).create(vae)\n\n  for epoch in range(FLAGS.num_epochs):\n    for batch in tfds.as_numpy(train_ds):\n      batch['image'] = batch['image'].reshape(-1, 784) / 255.0\n      optimizer = train_step(optimizer, batch)\n\n    z = np.random.normal(size=(64, 20))\n    metrics, comparison, sample = eval(optimizer.target, test_ds, z)\n    save_image(comparison, 'results/reconstruction_' + str(epoch) + '.png', nrow=8)\n    save_image(sample, 'results/sample_' + str(epoch) + '.png', nrow=8)\n\n    print(\"eval epoch: {}, loss: {:.4f}, BCE: {:.4f}, KLD: {:.4f}\".format(\n      epoch + 1, metrics['loss'], metrics['bce'], metrics['kld']\n    ))\n\n\nif __name__ == '__main__':\n    app.run(main)\n", "path": "examples/vae/main.py"}]}