Trying to implement a simple regime switch model in Edward

edward_regime_switch.py

import edward as ed
from edward import models as md
import tensorflow as tf
import pandas as pd
import numpy as np

# Using float32's everywhere because there's no DiscreteUniform, so typing is weird.
data = pd.read_csv('txtdata.csv', header=None).values[:,0].astype(dtype=np.float32)

n = len(data)
alpha_global = 1./n

# Model
lambda_1 = md.Exponential(alpha_global)
lambda_2 = md.Exponential(alpha_global)
tau = tf.nn.softplus(tf.Variable(n / 2.))

lambda_ = tf.where(tf.cast(tf.range(n), tf.float32) < tau, tf.ones(n)*lambda_1, tf.ones(n)*lambda_2)
x = md.Poisson(lam=lambda_, value=tf.ones(n)*alpha_global)

# Inference
ql_alpha_1 = tf.nn.softplus(tf.Variable(tf.random_normal([])))
ql_alpha_2 = tf.nn.softplus(tf.Variable(tf.random_normal([])))

ql_1 = md.Exponential(ql_alpha_1)
ql_2 = md.Exponential(ql_alpha_2)

inference = ed.KLqp({lambda_1: ql_1, lambda_2: ql_2}, {x: data})

inference.run(n_iter=20)

# I get some weird results for the loss:
# Iteration  1 [  5%]: Loss = -231632961261790988410899370537385984.000
# Iteration  2 [ 10%]: Loss = 0.000
# Iteration  4 [ 20%]: Loss = -158650697632963589784083628032.000
# Iteration  6 [ 30%]: Loss = 0.000
# Iteration  8 [ 40%]: Loss = 0.000
# Iteration 10 [ 50%]: Loss = -158650697632963589784083628032.000
# Iteration 12 [ 60%]: Loss = 0.000
# Iteration 14 [ 70%]: Loss = 0.000
# Iteration 16 [ 80%]: Loss = -158650697632963589784083628032.000
# Iteration 18 [ 90%]: Loss = 0.000
# Iteration 20 [100%]: Loss = 0.000