[pull] master from altosaar:master

.env

@@ -0,0 +1,7 @@
+# dev.env - development configuration
+
+# suppress warnings for jax
+JAX_PLATFORM_NAME=cpu
+
+# suppress tensorflow warnings
+TF_CPP_MIN_LOG_LEVEL=2

.gitignore

@@ -0,0 +1,4 @@
+*.pyc
+launch.json
+settings.json
+*.code-workspace

README.md

@@ -34,29 +34,12 @@ step:   20000   train elbo: -101.51
 step:   20000           valid elbo: -105.02     valid log p(x): -99.11
 step:   30000   train elbo: -98.70
 step:   30000           valid elbo: -103.76     valid log p(x): -97.71
-step:   40000   train elbo: -104.31
-step:   40000           valid elbo: -103.71     valid log p(x): -97.27
-step:   50000   train elbo: -97.20
-step:   50000           valid elbo: -102.97     valid log p(x): -96.60
-step:   60000   train elbo: -97.50
-step:   60000           valid elbo: -102.82     valid log p(x): -96.49
-step:   70000   train elbo: -94.68
-step:   70000           valid elbo: -102.63     valid log p(x): -96.22
-step:   80000   train elbo: -92.86
-step:   80000           valid elbo: -102.53     valid log p(x): -96.09
-step:   90000   train elbo: -93.83
-step:   90000           valid elbo: -102.33     valid log p(x): -96.00
-step:   100000  train elbo: -93.91
-step:   100000          valid elbo: -102.48     valid log p(x): -95.92
-step:   110000  train elbo: -94.34
-step:   110000          valid elbo: -102.81     valid log p(x): -96.09
-step:   120000  train elbo: -88.63
-step:   120000          valid elbo: -102.53     valid log p(x): -95.80
-step:   130000  train elbo: -96.61
-step:   130000          valid elbo: -103.56     valid log p(x): -96.26
-step:   140000  train elbo: -94.92
-step:   140000          valid elbo: -102.81     valid log p(x): -95.86
-step:   150000  train elbo: -97.84
-step:   150000          valid elbo: -103.06     valid log p(x): -95.92
-step:   150000          test elbo: -101.64      test log p(x): -95.33
 ```
+
+Using jax:
+```
+Step 0         	Validation ELBO estimate: -507.485	Validation log p(x) estimate: -507.485
+Step 10000     	Validation ELBO estimate: -152.695	Validation log p(x) estimate: -152.695
+Step 20000     	Validation ELBO estimate: -150.413	Validation log p(x) estimate: -150.413
+Step 30000     	Validation ELBO estimate: -150.529	Validation log p(x) estimate: -150.529
+```

environment_jax.yml

@@ -0,0 +1,81 @@
+name: jax
+channels:
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - ca-certificates=2021.4.13=h06a4308_1
+  - certifi=2020.12.5=py39h06a4308_0
+  - ld_impl_linux-64=2.33.1=h53a641e_7
+  - libffi=3.3=he6710b0_2
+  - libgcc-ng=9.1.0=hdf63c60_0
+  - libstdcxx-ng=9.1.0=hdf63c60_0
+  - ncurses=6.2=he6710b0_1
+  - openssl=1.1.1k=h27cfd23_0
+  - pip=21.1.1=py39h06a4308_0
+  - python=3.9.5=hdb3f193_3
+  - readline=8.1=h27cfd23_0
+  - setuptools=52.0.0=py39h06a4308_0
+  - sqlite=3.35.4=hdfb4753_0
+  - tk=8.6.10=hbc83047_0
+  - tzdata=2020f=h52ac0ba_0
+  - wheel=0.36.2=pyhd3eb1b0_0
+  - xz=5.2.5=h7b6447c_0
+  - zlib=1.2.11=h7b6447c_3
+  - pip:
+    - absl-py==0.12.0
+    - astunparse==1.6.3
+    - attrs==21.2.0
+    - cachetools==4.2.2
+    - chardet==4.0.0
+    - chex==0.0.7
+    - cloudpickle==1.6.0
+    - decorator==5.0.9
+    - dill==0.3.3
+    - dm-haiku==0.0.5.dev0
+    - dm-tree==0.1.6
+    - flatbuffers==1.12
+    - future==0.18.2
+    - gast==0.4.0
+    - google-auth==1.30.0
+    - google-auth-oauthlib==0.4.4
+    - google-pasta==0.2.0
+    - googleapis-common-protos==1.53.0
+    - grpcio==1.34.1
+    - h5py==3.1.0
+    - idna==2.10
+    - jax==0.2.13
+    - jaxlib==0.1.67
+    - jmp==0.0.2
+    - keras-nightly==2.5.0.dev2021032900
+    - keras-preprocessing==1.1.2
+    - markdown==3.3.4
+    - numpy==1.19.5
+    - oauthlib==3.1.0
+    - opt-einsum==3.3.0
+    - optax==0.0.7
+    - promise==2.3
+    - protobuf==3.17.0
+    - pyasn1==0.4.8
+    - pyasn1-modules==0.2.8
+    - requests==2.25.1
+    - requests-oauthlib==1.3.0
+    - rsa==4.7.2
+    - scipy==1.6.3
+    - six==1.15.0
+    - tabulate==0.8.9
+    - tensorboard==2.5.0
+    - tensorboard-data-server==0.6.1
+    - tensorboard-plugin-wit==1.8.0
+    - tensorflow==2.5.0
+    - tensorflow-datasets==4.3.0
+    - tensorflow-estimator==2.5.0
+    - tensorflow-metadata==1.0.0
+    - termcolor==1.1.0
+    - tfp-nightly==0.14.0.dev20210521
+    - toolz==0.11.1
+    - tqdm==4.60.0
+    - typing-extensions==3.7.4.3
+    - urllib3==1.26.4
+    - werkzeug==2.0.1
+    - wrapt==1.12.1
+prefix: /home/jaan/miniconda3/envs/jax

setup.cfg

@@ -0,0 +1,2 @@
+[flake8]
+max-line-length = 88