Saving edward models as SavedModel for serving doesn't work

[ShantanuKumar]

I am trying to save edward model as SavedModel which works.

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import edward as ed
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
import pandas as pd
import json

from edward.models import Normal
from tensorflow.python.framework import ops
from tensorflow.python.saved_model import builder
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import signature_def_utils
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.util.tf_export import tf_export

plt.style.use('ggplot')
ed.set_seed(42)

def save_model(session, export_dir, inputs, outputs):
    # delete older export dir
    if tf.gfile.Exists(export_dir):
        tf.logging.warning("deleting export dir")
        tf.gfile.DeleteRecursively(export_dir)

    signature_def_map = {
        signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
        signature_def_utils.predict_signature_def(inputs, outputs)
    }
    b = builder.SavedModelBuilder(export_dir)
    b.add_meta_graph_and_variables(
        session,
        tags=[tag_constants.SERVING],
        signature_def_map=signature_def_map,
        assets_collection=ops.get_collection(ops.GraphKeys.ASSET_FILEPATHS),
        legacy_init_op=None,
        clear_devices=True)
    b.save()

def build_toy_dataset(N, w):
    D = len(w)
    x = np.random.normal(0.0, 2.0, size=(N, D))
    y = np.dot(x, w) + np.random.normal(0.0, 0.01, size=N)
    return x, y

# Data
N = 40  # number of data points
D = 10  # number of features

w_true = np.random.randn(D) * 0.5
x_train, y_train = build_toy_dataset(N, w_true)
x_test, y_test = build_toy_dataset(N, w_true)

# Model
X = tf.placeholder(tf.float32, [None, D])
w = Normal(loc=tf.zeros(D), scale=tf.ones(D))
b = Normal(loc=tf.zeros(1), scale=tf.ones(1))
y = Normal(loc=ed.dot(X, w) + b, scale=tf.ones(1))

# Inference
qw = Normal(loc=tf.get_variable("qw/loc", [D]),
            scale=tf.nn.softplus(tf.get_variable("qw/scale", [D])))
qb = Normal(loc=tf.get_variable("qb/loc", [1]),
            scale=tf.nn.softplus(tf.get_variable("qb/scale", [1])))
inference = ed.KLqp({w: qw, b: qb}, data={X: x_train, y: y_train})
inference.run(n_samples=5, n_iter=250)

# Criticism
y_post = ed.copy(y, {w: qw, b: qb})

# save model
save_model(ed.get_session(), "./saved_model",
           inputs={"X": X}, outputs={"y": y_post})

But when I try to get a point prediction using saved_model_cli, it throws error.

saved_model_cli run --dir saved_model --tag_set serve  --signature_def serving_default --input_exprs 'X=np.array([[1., 2., 3., 4., 5., 6., 7., 8., 9., 10.]])'

2018-04-18 18:22:05.831559: I tensorflow/core/platform/cpu_feature_guard.cc:137] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.1 SSE4.2 AVX AVX2 FMA
Traceback (most recent call last):
  File "/usr/local/bin/saved_model_cli", line 11, in <module>
    sys.exit(main())
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/tools/saved_model_cli.py", line 648, in main
    args.func(args)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/tools/saved_model_cli.py", line 526, in run
    args.overwrite, tf_debug=args.tf_debug)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/tools/saved_model_cli.py", line 295, in run_saved_model_with_feed_dict
    outputs = sess.run(output_tensor_names_sorted, feed_dict=inputs_feed_dict)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 889, in run
    run_metadata_ptr)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 1105, in _run
    self._graph, fetches, feed_dict_tensor, feed_handles=feed_handles)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 414, in __init__
    self._fetch_mapper = _FetchMapper.for_fetch(fetches)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 234, in for_fetch
    return _ListFetchMapper(fetch)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 341, in __init__
    self._mappers = [_FetchMapper.for_fetch(fetch) for fetch in fetches]
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 242, in for_fetch
    return _ElementFetchMapper(fetches, contraction_fn)
  File "/usr/local/lib/python2.7/dist-packages/tensorflow/python/client/session.py", line 281, in __init__
    'Tensor. (%s)' % (fetch, str(e)))
ValueError: Fetch argument u'copied/Normal_2/' cannot be interpreted as a Tensor. ("The name 'copied/Normal_2/' refers to an Operation not in the graph.")