How to run/use bayesian_neural_network example

Hello all,

I found in the examples bayesian_neural_network.py which should train a bayesian neural network to classify MNIST digits.
Please is there any blog, readme or notes on how to use this script - it would be helpful to understand more.
Thank you.

it fails to run it on Window10 machine

>python bayesian_neural_network.py
2020-04-21 15:46:55.542625: I tensorflow/stream_executor/platform/default/dso_loader.cc:44] Successfully opened dynamic library cudart64_101.dll
Traceback (most recent call last):
  File "bayesian_neural_network.py", line 366, in <module>
    app.run(main)
  File "C:\Anaconda3\envs\tf-gpu\lib\site-packages\absl\app.py", line 299, in run
    _run_main(main, args)
  File "C:\Anaconda3\envs\tf-gpu\lib\site-packages\absl\app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "bayesian_neural_network.py", line 297, in main
    tf.io.gfile.makedirs(FLAGS.model_dir)
  File "C:\Anaconda3\envs\tf-gpu\lib\site-packages\tensorflow_core\python\lib\io\file_io.py", line 455, in recursive_create_dir_v2
    pywrap_tensorflow.RecursivelyCreateDir(compat.as_bytes(path))
tensorflow.python.framework.errors_impl.NotFoundError: Failed to create a directory: /tmp\bayesian_neural_network; No such file or directory

On my Linux machine, I got following errors

2020-04-21 16:20:12.937404: E tensorflow/stream_executor/cuda/cuda_dnn.cc:329] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2020-04-21 16:20:13.198242: E tensorflow/stream_executor/cuda/cuda_dnn.cc:329] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2020-04-21 16:20:13.198916: W tensorflow/core/common_runtime/base_collective_executor.cc:217] BaseCollectiveExecutor::StartAbort Unknown: Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.
         [[{{node sequential/conv2d_flipout/Conv2D}}]]
2020-04-21 16:20:13.343447: E tensorflow/stream_executor/cuda/cuda_dnn.cc:329] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
2020-04-21 16:20:13.348532: E tensorflow/stream_executor/cuda/cuda_dnn.cc:329] Could not create cudnn handle: CUDNN_STATUS_INTERNAL_ERROR
Traceback (most recent call last):
  File "bayesian_neural_network.py", line 366, in <module>
    app.run(main)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/absl/app.py", line 299, in run
    _run_main(main, args)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/absl/app.py", line 250, in _run_main
    sys.exit(main(argv))
  File "bayesian_neural_network.py", line 319, in main
    batch_x, batch_y)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/training.py", line 1123, in train_on_batch
    outputs = self.train_function(ins)  # pylint: disable=not-callable
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/keras/backend.py", line 3727, in __call__
    outputs = self._graph_fn(*converted_inputs)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1551, in __call__
    return self._call_impl(args, kwargs)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1591, in _call_impl
    return self._call_flat(args, self.captured_inputs, cancellation_manager)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 1692, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/eager/function.py", line 545, in call
    ctx=ctx)
  File "/home/anaconda3/envs/xAI/lib/python3.7/site-packages/tensorflow_core/python/eager/execute.py", line 67, in quick_execute
    six.raise_from(core._status_to_exception(e.code, message), None)
  File "<string>", line 3, in raise_from
tensorflow.python.framework.errors_impl.UnknownError:  Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.
         [[node sequential/conv2d_flipout/Conv2D (defined at bayesian_neural_network.py:319) ]] [Op:__inference_keras_scratch_graph_3715]

Function call stack:
keras_scratch_graph

Start with examples/bayesian_neural_network.py, especially the paths around lines 297 and 319 shown in the reports, and review how the example is invoked on Windows and Linux. Document the prerequisites, usage steps, and platform-specific failure guidance; done means a newcomer can run the MNIST example or identify the relevant environment issue.