We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
Allow testing of examples not parameterized by 'num_epochs'
The AIR example isn't run by the inference testing machinery, as it doesn't have a `--num-epochs` argument.

I could add such an argument, but this isn't a very natural parameterization for this example. Also, running an entire epoch for each example seems likely to be overkill in some cases, including AIR. (Assuming an epoch is a pass through the whole dataset.)

One idea is that the test should look for examples with a dedicated`--ci-test` argument or similar. This is a bit more flexible, though we might pay for this with a little extra code in each example.

--- 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/air/main.py`
Content:
```
1 """
2 This attempts (currently unsuccessfully) to reproduce some of the
3 results on the multi-mnist data set described in [1].
4 
5 [1] Eslami, SM Ali, et al. "Attend, infer, repeat: Fast scene
6 understanding with generative models." Advances in Neural Information
7 Processing Systems. 2016.
8 """
9 
10 import os
11 import time
12 import argparse
13 from functools import partial
14 import numpy as np
15 
16 import torch
17 from torch.autograd import Variable
18 
19 import pyro
20 import pyro.optim as optim
21 import pyro.poutine as poutine
22 from pyro.infer import SVI
23 
24 import visdom
25 
26 from air import AIR
27 from viz import draw_many, post_process_latents
28 
29 parser = argparse.ArgumentParser(description="Pyro AIR example", argument_default=argparse.SUPPRESS)
30 parser.add_argument('-n', '--num-steps', type=int, default=int(1e8),
31                     help='number of optimization steps to take')
32 parser.add_argument('-b', '--batch-size', type=int, default=64,
33                     help='batch size')
34 parser.add_argument('--progress-every', type=int, default=1,
35                     help='number of steps between writing progress to stdout')
36 parser.add_argument('--baseline-scalar', type=float,
37                     help='scale the output of the baseline nets by this value')
38 parser.add_argument('--no-baselines', action='store_true', default=False,
39                     help='do not use data dependent baselines')
40 parser.add_argument('--encoder-net', type=int, nargs='+', default=[200],
41                     help='encoder net hidden layer sizes')
42 parser.add_argument('--decoder-net', type=int, nargs='+', default=[200],
43                     help='decoder net hidden layer sizes')
44 parser.add_argument('--predict-net', type=int, nargs='+',
45                     help='predict net hidden layer sizes')
46 parser.add_argument('--embed-net', type=int, nargs='+',
47                     help='embed net architecture')
48 parser.add_argument('--bl-predict-net', type=int, nargs='+',
49                     help='baseline predict net hidden layer sizes')
50 parser.add_argument('--non-linearity', type=str,
51                     help='non linearity to use throughout')
52 parser.add_argument('--viz', action='store_true', default=False,
53                     help='generate vizualizations during optimization')
54 parser.add_argument('--viz-every', type=int, default=100,
55                     help='number of steps between vizualizations')
56 parser.add_argument('--visdom-env', default='main',
57                     help='visdom enviroment name')
58 # parser.add_argument('--checkpoint', action='store_true',
59 #                     help='periodically persist parameters')
60 # parser.add_argument('--checkpoint-every', type=int, default=1000,
61 #                     help='number of steps between checkpoints')
62 parser.add_argument('--cuda', action='store_true', default=False,
63                     help='use cuda')
64 parser.add_argument('-t', '--model-steps', type=int, default=3,
65                     help='number of time steps')
66 parser.add_argument('--rnn-hidden-size', type=int, default=256,
67                     help='rnn hidden size')
68 parser.add_argument('--encoder-latent-size', type=int, default=50,
69                     help='attention window encoder/decoder latent space size')
70 parser.add_argument('--decoder-output-bias', type=float,
71                     help='bias added to decoder output (prior to applying non-linearity)')
72 parser.add_argument('--window-size', type=int, default=28,
73                     help='attention window size')
74 parser.add_argument('--z-pres-prior', type=float, default=None,
75                     help='prior success probability for z_pres')
76 parser.add_argument('--anneal-prior', action='store_true', default=False,
77                     help='anneal z_pres prior during optimization')
78 parser.add_argument('--anneal-prior-from', type=float, default=0.99,
79                     help='initial z_pres prior prob')
80 parser.add_argument('--anneal-prior-over', type=int, default=100000,
81                     help='number of steps over which to anneal the prior')
82 parser.add_argument('--no-masking', action='store_true', default=False,
83                     help='do not mask out the costs of unused choices')
84 parser.add_argument('--fudge-z-pres', action='store_true', default=False,
85                     help='fudge z_pres to remove discreteness for testing')
86 parser.add_argument('--seed', type=int, help='random seed', default=None)
87 parser.add_argument('--print-modules', action='store_true',
88                     help='write the network architecture to stdout')
89 
90 args = parser.parse_args()
91 # print(args)
92 
93 if args.seed is not None:
94     pyro.set_rng_seed(args.seed)
95 
96 # Load data.
97 infile = './data/multi_mnist_train_uint8.npz'
98 if not os.path.exists(infile):
99     print('Could not find the dataset at {}'.format(infile))
100     print('Run "python multi_mnist.py" to generate it.')
101     exit()
102 X_np = np.load(infile)['x'].astype(np.float32)
103 X_np /= 255.0
104 X = Variable(torch.from_numpy(X_np))
105 X_size = X.size(0)
106 if args.cuda:
107     X = X.cuda()
108 
109 
110 # Yields the following distribution over the number of steps (when
111 # taking a maximum of 3 steps):
112 # p(0) = 0.4
113 # p(1) = 0.3
114 # p(2) = 0.2
115 # p(3) = 0.1
116 def default_z_pres_prior_p(t):
117     if t == 0:
118         return 0.6
119     elif t == 1:
120         return 0.5
121     else:
122         return 0.33
123 
124 
125 # Implements "prior annealing" as described in this blog post:
126 # http://akosiorek.github.io/ml/2017/09/03/implementing-air.html
127 def annealed_z_pres_prior_p(opt_step, time_step):
128     p_0 = args.anneal_prior_from
129     p_final = args.z_pres_prior or default_z_pres_prior_p(time_step)
130     s = min(opt_step / args.anneal_prior_over, 1.0)
131     return p_final * s + p_0 * (1 - s)
132 
133 
134 def z_pres_prior_p(opt_step, time_step):
135     if args.anneal_prior:
136         return annealed_z_pres_prior_p(opt_step, time_step)
137     elif args.z_pres_prior:
138         return args.z_pres_prior
139     else:
140         return default_z_pres_prior_p(time_step)
141 
142 
143 model_arg_keys = ['window_size',
144                   'rnn_hidden_size',
145                   'decoder_output_bias',
146                   'baseline_scalar',
147                   'encoder_net',
148                   'decoder_net',
149                   'predict_net',
150                   'embed_net',
151                   'bl_predict_net',
152                   'non_linearity',
153                   'fudge_z_pres',
154                   'print_modules']
155 model_args = {key: getattr(args, key) for key in model_arg_keys if key in args}
156 air = AIR(
157     num_steps=args.model_steps,
158     x_size=50,
159     use_masking=not args.no_masking,
160     use_baselines=not args.no_baselines,
161     z_what_size=args.encoder_latent_size,
162     use_cuda=args.cuda,
163     **model_args
164 )
165 
166 vis = visdom.Visdom(env=args.visdom_env)
167 # Viz sample from prior.
168 if args.viz:
169     z, x = air.prior(5, z_pres_prior_p=partial(z_pres_prior_p, 0))
170     vis.images(draw_many(x, post_process_latents(z)))
171 
172 t0 = time.time()
173 examples_to_viz = X[9:14]
174 
175 
176 # Do inference.
177 def per_param_optim_args(module_name, param_name, tags):
178     lr = 1e-3 if 'baseline' in tags else 1e-4
179     return {'lr': lr}
180 
181 
182 svi = SVI(air.model, air.guide,
183           optim.Adam(per_param_optim_args),
184           loss='ELBO',
185           trace_graph=True)
186 
187 for i in range(args.num_steps):
188 
189     loss = svi.step(X, args.batch_size, z_pres_prior_p=partial(z_pres_prior_p, i))
190 
191     if i % args.progress_every == 0:
192         print('i={}, epochs={:.2f}, elapsed={:.2f}, elbo={:.2f}'.format(
193             i,
194             (i * args.batch_size) / X_size,
195             (time.time() - t0) / 3600,
196             loss / X_size))
197 
198     if args.viz and (i + 1) % args.viz_every == 0:
199         trace = poutine.trace(air.guide).get_trace(examples_to_viz, 0)
200         z, recons = poutine.replay(air.prior, trace)(examples_to_viz.size(0))
201         z_wheres = post_process_latents(z)
202 
203         # Show data with inferred objection positions.
204         vis.images(draw_many(examples_to_viz, z_wheres))
205         # Show reconstructions of data.
206         vis.images(draw_many(recons, z_wheres))
207 
208         # TODO: Report accuracy on predictions of object counts.
209
```
--- 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/air/main.py b/examples/air/main.py
--- a/examples/air/main.py
+++ b/examples/air/main.py
@@ -8,9 +8,11 @@
 """
 
 import os
+import sys
 import time
 import argparse
 from functools import partial
+from subprocess import check_call
 import numpy as np
 
 import torch
@@ -96,9 +98,10 @@
 # Load data.
 infile = './data/multi_mnist_train_uint8.npz'
 if not os.path.exists(infile):
-    print('Could not find the dataset at {}'.format(infile))
-    print('Run "python multi_mnist.py" to generate it.')
-    exit()
+    print('Running multi_mnist.py to generate dataset at {}...'.format(infile))
+    multi_mnist_py = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'multi_mnist.py')
+    check_call([sys.executable, multi_mnist_py])
+    print('Finished running multi_mnist.py.')
 X_np = np.load(infile)['x'].astype(np.float32)
 X_np /= 255.0
 X = Variable(torch.from_numpy(X_np))

{"golden_diff": "diff --git a/examples/air/main.py b/examples/air/main.py\n--- a/examples/air/main.py\n+++ b/examples/air/main.py\n@@ -8,9 +8,11 @@\n \"\"\"\n \n import os\n+import sys\n import time\n import argparse\n from functools import partial\n+from subprocess import check_call\n import numpy as np\n \n import torch\n@@ -96,9 +98,10 @@\n # Load data.\n infile = './data/multi_mnist_train_uint8.npz'\n if not os.path.exists(infile):\n-    print('Could not find the dataset at {}'.format(infile))\n-    print('Run \"python multi_mnist.py\" to generate it.')\n-    exit()\n+    print('Running multi_mnist.py to generate dataset at {}...'.format(infile))\n+    multi_mnist_py = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'multi_mnist.py')\n+    check_call([sys.executable, multi_mnist_py])\n+    print('Finished running multi_mnist.py.')\n X_np = np.load(infile)['x'].astype(np.float32)\n X_np /= 255.0\n X = Variable(torch.from_numpy(X_np))\n", "issue": "Allow testing of examples not parameterized by 'num_epochs'\nThe AIR example isn't run by the inference testing machinery, as it doesn't have a `--num-epochs` argument.\r\n\r\nI could add such an argument, but this isn't a very natural parameterization for this example. Also, running an entire epoch for each example seems likely to be overkill in some cases, including AIR. (Assuming an epoch is a pass through the whole dataset.)\r\n\r\nOne idea is that the test should look for examples with a dedicated`--ci-test` argument or similar. This is a bit more flexible, though we might pay for this with a little extra code in each example.\n", "before_files": [{"content": "\"\"\"\nThis attempts (currently unsuccessfully) to reproduce some of the\nresults on the multi-mnist data set described in [1].\n\n[1] Eslami, SM Ali, et al. \"Attend, infer, repeat: Fast scene\nunderstanding with generative models.\" Advances in Neural Information\nProcessing Systems. 2016.\n\"\"\"\n\nimport os\nimport time\nimport argparse\nfrom functools import partial\nimport numpy as np\n\nimport torch\nfrom torch.autograd import Variable\n\nimport pyro\nimport pyro.optim as optim\nimport pyro.poutine as poutine\nfrom pyro.infer import SVI\n\nimport visdom\n\nfrom air import AIR\nfrom viz import draw_many, post_process_latents\n\nparser = argparse.ArgumentParser(description=\"Pyro AIR example\", argument_default=argparse.SUPPRESS)\nparser.add_argument('-n', '--num-steps', type=int, default=int(1e8),\n                    help='number of optimization steps to take')\nparser.add_argument('-b', '--batch-size', type=int, default=64,\n                    help='batch size')\nparser.add_argument('--progress-every', type=int, default=1,\n                    help='number of steps between writing progress to stdout')\nparser.add_argument('--baseline-scalar', type=float,\n                    help='scale the output of the baseline nets by this value')\nparser.add_argument('--no-baselines', action='store_true', default=False,\n                    help='do not use data dependent baselines')\nparser.add_argument('--encoder-net', type=int, nargs='+', default=[200],\n                    help='encoder net hidden layer sizes')\nparser.add_argument('--decoder-net', type=int, nargs='+', default=[200],\n                    help='decoder net hidden layer sizes')\nparser.add_argument('--predict-net', type=int, nargs='+',\n                    help='predict net hidden layer sizes')\nparser.add_argument('--embed-net', type=int, nargs='+',\n                    help='embed net architecture')\nparser.add_argument('--bl-predict-net', type=int, nargs='+',\n                    help='baseline predict net hidden layer sizes')\nparser.add_argument('--non-linearity', type=str,\n                    help='non linearity to use throughout')\nparser.add_argument('--viz', action='store_true', default=False,\n                    help='generate vizualizations during optimization')\nparser.add_argument('--viz-every', type=int, default=100,\n                    help='number of steps between vizualizations')\nparser.add_argument('--visdom-env', default='main',\n                    help='visdom enviroment name')\n# parser.add_argument('--checkpoint', action='store_true',\n#                     help='periodically persist parameters')\n# parser.add_argument('--checkpoint-every', type=int, default=1000,\n#                     help='number of steps between checkpoints')\nparser.add_argument('--cuda', action='store_true', default=False,\n                    help='use cuda')\nparser.add_argument('-t', '--model-steps', type=int, default=3,\n                    help='number of time steps')\nparser.add_argument('--rnn-hidden-size', type=int, default=256,\n                    help='rnn hidden size')\nparser.add_argument('--encoder-latent-size', type=int, default=50,\n                    help='attention window encoder/decoder latent space size')\nparser.add_argument('--decoder-output-bias', type=float,\n                    help='bias added to decoder output (prior to applying non-linearity)')\nparser.add_argument('--window-size', type=int, default=28,\n                    help='attention window size')\nparser.add_argument('--z-pres-prior', type=float, default=None,\n                    help='prior success probability for z_pres')\nparser.add_argument('--anneal-prior', action='store_true', default=False,\n                    help='anneal z_pres prior during optimization')\nparser.add_argument('--anneal-prior-from', type=float, default=0.99,\n                    help='initial z_pres prior prob')\nparser.add_argument('--anneal-prior-over', type=int, default=100000,\n                    help='number of steps over which to anneal the prior')\nparser.add_argument('--no-masking', action='store_true', default=False,\n                    help='do not mask out the costs of unused choices')\nparser.add_argument('--fudge-z-pres', action='store_true', default=False,\n                    help='fudge z_pres to remove discreteness for testing')\nparser.add_argument('--seed', type=int, help='random seed', default=None)\nparser.add_argument('--print-modules', action='store_true',\n                    help='write the network architecture to stdout')\n\nargs = parser.parse_args()\n# print(args)\n\nif args.seed is not None:\n    pyro.set_rng_seed(args.seed)\n\n# Load data.\ninfile = './data/multi_mnist_train_uint8.npz'\nif not os.path.exists(infile):\n    print('Could not find the dataset at {}'.format(infile))\n    print('Run \"python multi_mnist.py\" to generate it.')\n    exit()\nX_np = np.load(infile)['x'].astype(np.float32)\nX_np /= 255.0\nX = Variable(torch.from_numpy(X_np))\nX_size = X.size(0)\nif args.cuda:\n    X = X.cuda()\n\n\n# Yields the following distribution over the number of steps (when\n# taking a maximum of 3 steps):\n# p(0) = 0.4\n# p(1) = 0.3\n# p(2) = 0.2\n# p(3) = 0.1\ndef default_z_pres_prior_p(t):\n    if t == 0:\n        return 0.6\n    elif t == 1:\n        return 0.5\n    else:\n        return 0.33\n\n\n# Implements \"prior annealing\" as described in this blog post:\n# http://akosiorek.github.io/ml/2017/09/03/implementing-air.html\ndef annealed_z_pres_prior_p(opt_step, time_step):\n    p_0 = args.anneal_prior_from\n    p_final = args.z_pres_prior or default_z_pres_prior_p(time_step)\n    s = min(opt_step / args.anneal_prior_over, 1.0)\n    return p_final * s + p_0 * (1 - s)\n\n\ndef z_pres_prior_p(opt_step, time_step):\n    if args.anneal_prior:\n        return annealed_z_pres_prior_p(opt_step, time_step)\n    elif args.z_pres_prior:\n        return args.z_pres_prior\n    else:\n        return default_z_pres_prior_p(time_step)\n\n\nmodel_arg_keys = ['window_size',\n                  'rnn_hidden_size',\n                  'decoder_output_bias',\n                  'baseline_scalar',\n                  'encoder_net',\n                  'decoder_net',\n                  'predict_net',\n                  'embed_net',\n                  'bl_predict_net',\n                  'non_linearity',\n                  'fudge_z_pres',\n                  'print_modules']\nmodel_args = {key: getattr(args, key) for key in model_arg_keys if key in args}\nair = AIR(\n    num_steps=args.model_steps,\n    x_size=50,\n    use_masking=not args.no_masking,\n    use_baselines=not args.no_baselines,\n    z_what_size=args.encoder_latent_size,\n    use_cuda=args.cuda,\n    **model_args\n)\n\nvis = visdom.Visdom(env=args.visdom_env)\n# Viz sample from prior.\nif args.viz:\n    z, x = air.prior(5, z_pres_prior_p=partial(z_pres_prior_p, 0))\n    vis.images(draw_many(x, post_process_latents(z)))\n\nt0 = time.time()\nexamples_to_viz = X[9:14]\n\n\n# Do inference.\ndef per_param_optim_args(module_name, param_name, tags):\n    lr = 1e-3 if 'baseline' in tags else 1e-4\n    return {'lr': lr}\n\n\nsvi = SVI(air.model, air.guide,\n          optim.Adam(per_param_optim_args),\n          loss='ELBO',\n          trace_graph=True)\n\nfor i in range(args.num_steps):\n\n    loss = svi.step(X, args.batch_size, z_pres_prior_p=partial(z_pres_prior_p, i))\n\n    if i % args.progress_every == 0:\n        print('i={}, epochs={:.2f}, elapsed={:.2f}, elbo={:.2f}'.format(\n            i,\n            (i * args.batch_size) / X_size,\n            (time.time() - t0) / 3600,\n            loss / X_size))\n\n    if args.viz and (i + 1) % args.viz_every == 0:\n        trace = poutine.trace(air.guide).get_trace(examples_to_viz, 0)\n        z, recons = poutine.replay(air.prior, trace)(examples_to_viz.size(0))\n        z_wheres = post_process_latents(z)\n\n        # Show data with inferred objection positions.\n        vis.images(draw_many(examples_to_viz, z_wheres))\n        # Show reconstructions of data.\n        vis.images(draw_many(recons, z_wheres))\n\n        # TODO: Report accuracy on predictions of object counts.\n", "path": "examples/air/main.py"}], "after_files": [{"content": "\"\"\"\nThis attempts (currently unsuccessfully) to reproduce some of the\nresults on the multi-mnist data set described in [1].\n\n[1] Eslami, SM Ali, et al. \"Attend, infer, repeat: Fast scene\nunderstanding with generative models.\" Advances in Neural Information\nProcessing Systems. 2016.\n\"\"\"\n\nimport os\nimport sys\nimport time\nimport argparse\nfrom functools import partial\nfrom subprocess import check_call\nimport numpy as np\n\nimport torch\nfrom torch.autograd import Variable\n\nimport pyro\nimport pyro.optim as optim\nimport pyro.poutine as poutine\nfrom pyro.infer import SVI\n\nimport visdom\n\nfrom air import AIR\nfrom viz import draw_many, post_process_latents\n\nparser = argparse.ArgumentParser(description=\"Pyro AIR example\", argument_default=argparse.SUPPRESS)\nparser.add_argument('-n', '--num-steps', type=int, default=int(1e8),\n                    help='number of optimization steps to take')\nparser.add_argument('-b', '--batch-size', type=int, default=64,\n                    help='batch size')\nparser.add_argument('--progress-every', type=int, default=1,\n                    help='number of steps between writing progress to stdout')\nparser.add_argument('--baseline-scalar', type=float,\n                    help='scale the output of the baseline nets by this value')\nparser.add_argument('--no-baselines', action='store_true', default=False,\n                    help='do not use data dependent baselines')\nparser.add_argument('--encoder-net', type=int, nargs='+', default=[200],\n                    help='encoder net hidden layer sizes')\nparser.add_argument('--decoder-net', type=int, nargs='+', default=[200],\n                    help='decoder net hidden layer sizes')\nparser.add_argument('--predict-net', type=int, nargs='+',\n                    help='predict net hidden layer sizes')\nparser.add_argument('--embed-net', type=int, nargs='+',\n                    help='embed net architecture')\nparser.add_argument('--bl-predict-net', type=int, nargs='+',\n                    help='baseline predict net hidden layer sizes')\nparser.add_argument('--non-linearity', type=str,\n                    help='non linearity to use throughout')\nparser.add_argument('--viz', action='store_true', default=False,\n                    help='generate vizualizations during optimization')\nparser.add_argument('--viz-every', type=int, default=100,\n                    help='number of steps between vizualizations')\nparser.add_argument('--visdom-env', default='main',\n                    help='visdom enviroment name')\n# parser.add_argument('--checkpoint', action='store_true',\n#                     help='periodically persist parameters')\n# parser.add_argument('--checkpoint-every', type=int, default=1000,\n#                     help='number of steps between checkpoints')\nparser.add_argument('--cuda', action='store_true', default=False,\n                    help='use cuda')\nparser.add_argument('-t', '--model-steps', type=int, default=3,\n                    help='number of time steps')\nparser.add_argument('--rnn-hidden-size', type=int, default=256,\n                    help='rnn hidden size')\nparser.add_argument('--encoder-latent-size', type=int, default=50,\n                    help='attention window encoder/decoder latent space size')\nparser.add_argument('--decoder-output-bias', type=float,\n                    help='bias added to decoder output (prior to applying non-linearity)')\nparser.add_argument('--window-size', type=int, default=28,\n                    help='attention window size')\nparser.add_argument('--z-pres-prior', type=float, default=None,\n                    help='prior success probability for z_pres')\nparser.add_argument('--anneal-prior', action='store_true', default=False,\n                    help='anneal z_pres prior during optimization')\nparser.add_argument('--anneal-prior-from', type=float, default=0.99,\n                    help='initial z_pres prior prob')\nparser.add_argument('--anneal-prior-over', type=int, default=100000,\n                    help='number of steps over which to anneal the prior')\nparser.add_argument('--no-masking', action='store_true', default=False,\n                    help='do not mask out the costs of unused choices')\nparser.add_argument('--fudge-z-pres', action='store_true', default=False,\n                    help='fudge z_pres to remove discreteness for testing')\nparser.add_argument('--seed', type=int, help='random seed', default=None)\nparser.add_argument('--print-modules', action='store_true',\n                    help='write the network architecture to stdout')\n\nargs = parser.parse_args()\n# print(args)\n\nif args.seed is not None:\n    pyro.set_rng_seed(args.seed)\n\n# Load data.\ninfile = './data/multi_mnist_train_uint8.npz'\nif not os.path.exists(infile):\n    print('Running multi_mnist.py to generate dataset at {}...'.format(infile))\n    multi_mnist_py = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'multi_mnist.py')\n    check_call([sys.executable, multi_mnist_py])\n    print('Finished running multi_mnist.py.')\nX_np = np.load(infile)['x'].astype(np.float32)\nX_np /= 255.0\nX = Variable(torch.from_numpy(X_np))\nX_size = X.size(0)\nif args.cuda:\n    X = X.cuda()\n\n\n# Yields the following distribution over the number of steps (when\n# taking a maximum of 3 steps):\n# p(0) = 0.4\n# p(1) = 0.3\n# p(2) = 0.2\n# p(3) = 0.1\ndef default_z_pres_prior_p(t):\n    if t == 0:\n        return 0.6\n    elif t == 1:\n        return 0.5\n    else:\n        return 0.33\n\n\n# Implements \"prior annealing\" as described in this blog post:\n# http://akosiorek.github.io/ml/2017/09/03/implementing-air.html\ndef annealed_z_pres_prior_p(opt_step, time_step):\n    p_0 = args.anneal_prior_from\n    p_final = args.z_pres_prior or default_z_pres_prior_p(time_step)\n    s = min(opt_step / args.anneal_prior_over, 1.0)\n    return p_final * s + p_0 * (1 - s)\n\n\ndef z_pres_prior_p(opt_step, time_step):\n    if args.anneal_prior:\n        return annealed_z_pres_prior_p(opt_step, time_step)\n    elif args.z_pres_prior:\n        return args.z_pres_prior\n    else:\n        return default_z_pres_prior_p(time_step)\n\n\nmodel_arg_keys = ['window_size',\n                  'rnn_hidden_size',\n                  'decoder_output_bias',\n                  'baseline_scalar',\n                  'encoder_net',\n                  'decoder_net',\n                  'predict_net',\n                  'embed_net',\n                  'bl_predict_net',\n                  'non_linearity',\n                  'fudge_z_pres',\n                  'print_modules']\nmodel_args = {key: getattr(args, key) for key in model_arg_keys if key in args}\nair = AIR(\n    num_steps=args.model_steps,\n    x_size=50,\n    use_masking=not args.no_masking,\n    use_baselines=not args.no_baselines,\n    z_what_size=args.encoder_latent_size,\n    use_cuda=args.cuda,\n    **model_args\n)\n\nvis = visdom.Visdom(env=args.visdom_env)\n# Viz sample from prior.\nif args.viz:\n    z, x = air.prior(5, z_pres_prior_p=partial(z_pres_prior_p, 0))\n    vis.images(draw_many(x, post_process_latents(z)))\n\nt0 = time.time()\nexamples_to_viz = X[9:14]\n\n\n# Do inference.\ndef per_param_optim_args(module_name, param_name, tags):\n    lr = 1e-3 if 'baseline' in tags else 1e-4\n    return {'lr': lr}\n\n\nsvi = SVI(air.model, air.guide,\n          optim.Adam(per_param_optim_args),\n          loss='ELBO',\n          trace_graph=True)\n\nfor i in range(args.num_steps):\n\n    loss = svi.step(X, args.batch_size, z_pres_prior_p=partial(z_pres_prior_p, i))\n\n    if i % args.progress_every == 0:\n        print('i={}, epochs={:.2f}, elapsed={:.2f}, elbo={:.2f}'.format(\n            i,\n            (i * args.batch_size) / X_size,\n            (time.time() - t0) / 3600,\n            loss / X_size))\n\n    if args.viz and (i + 1) % args.viz_every == 0:\n        trace = poutine.trace(air.guide).get_trace(examples_to_viz, 0)\n        z, recons = poutine.replay(air.prior, trace)(examples_to_viz.size(0))\n        z_wheres = post_process_latents(z)\n\n        # Show data with inferred objection positions.\n        vis.images(draw_many(examples_to_viz, z_wheres))\n        # Show reconstructions of data.\n        vis.images(draw_many(recons, z_wheres))\n\n        # TODO: Report accuracy on predictions of object counts.\n", "path": "examples/air/main.py"}]}