model_CVAE.py

#### MAIN
from itertools import combinations
import collections
import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp
from base import *

EPS = 1e-4

class CVAE(tf.keras.Model):

    def __init__(self, n_params):

        super(CVAE, self).__init__()

        self.n_params = n_params

        n_recog, n_scale, n_gener_z, n_gener_kappa = self.n_params['n_recog'], self.n_params['n_scale'], self.n_params['n_gener_z'], self.n_params['n_gener_kappa']
        n_recog_y, n_gener_y = self.n_params['n_recog_y'], self.n_params['n_gener_y']
        n_recog_kappa = self.n_params['n_recog_kappa']
        n_z, n_x, n_y, n_kappa = self.n_params['n_z'], self.n_params['n_x'], self.n_params['n_y'], self.n_params['n_kappa'],
        n_site = self.n_params['n_site']

        n_kappa_max = np.max(n_kappa)
        self.n_kappa_masks = tf.stack([tf.concat([tf.ones(n_kappa[site]), tf.zeros(n_kappa_max - n_kappa[site])], axis=-1) for site in range(n_site)])

        self.learning_rate = self.n_params['learning_rate']
        self.alpha_1 = self.n_params['alpha_1']

        # Define eta_varialbes in cpu

        with tf.device('/cpu:0'):

            self.recog_x_trans_mat_values = self.add_weight(name='recog_x_trans_mat_values', shape=(n_x**2, ), initializer=tf.constant_initializer((1 - EPS) * np.eye(n_x).flatten()))
            self.recog_x_trans_mean = self.add_weight(name='recog_x_trans_mean', shape=(n_x, ), initializer=tf.constant_initializer(np.zeros(n_x)))

            self.recog_x_init_prec_values = self.add_weight(name='recog_x_init_prec_values', shape=(n_x*(n_x+1)//2, ), initializer=tf.constant_initializer(np.zeros(n_x*(n_x+1)//2)))
            self.recog_x_init_mean = self.add_weight(name='recog_x_init_mean', shape=(n_x, ), initializer=tf.constant_initializer(np.zeros(n_x)))

            self.gener_x_trans_mat_values = self.add_weight(name='gener_x_trans_mat_values', shape=(n_x**2, ), initializer=tf.constant_initializer((1 - EPS) * np.eye(n_x).flatten()))
            self.gener_x_trans_mean = self.add_weight(name='gener_x_trans_mean', shape=(n_x, ), initializer=tf.constant_initializer(np.zeros(n_x)))

            self.gener_x_init_prec_values = self.add_weight(name='gener_x_init_prec_values', shape=(n_x*(n_x+1)//2, ), initializer=tf.constant_initializer(np.zeros(n_x*(n_x+1)//2)))
            self.gener_x_init_mean = self.add_weight(name='gener_x_init_mean', shape=(n_x, ), initializer=tf.constant_initializer(np.zeros(n_x)))

            self.gener_x_trans_prec_values =  self.add_weight(name='gener_x_affine_values', shape=(n_x*(n_x+1)//2, ), initializer=tf.constant_initializer(np.hstack([np.ones(n_x), np.zeros(n_x*(n_x+1)//2-n_x)])))


        # Define eta_varialbes in gpu
        self.x_affine_values =  self.add_weight(name='x_affine_values', shape=(n_x*(n_x+1)//2, ), initializer=tf.constant_initializer(np.hstack([np.ones(n_x) * tf.math.log(1 - (1 - EPS) ** 2), np.zeros(n_x*(n_x+1)//2-n_x)])))

        self.recog_kappa_g_mean = self.add_weight(name='recog_kappa_g_mean', shape=(n_x, ), initializer=tf.constant_initializer(np.zeros(n_x)))
        self.recog_kappa_g_prec_values = self.add_weight(name='recog_kappa_g_prec_values', shape=(n_x*(n_x+1)//2, ), initializer=tf.constant_initializer(np.zeros(n_x*(n_x+1)//2)))

        self.recog_y = tf.keras.models.Sequential([tf.keras.layers.Dense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_recog_y])
        self.recog_y_mean, self.recog_y_prec_values = tf.keras.Sequential([tf.keras.layers.Activation('leaky_relu'),  tf.keras.layers.Dense(n_x)]), tf.keras.Sequential([clip_by_value, TriangularDense(n_x)])

        self.recog_kappa = tf.keras.models.Sequential([ParallelDense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_recog_kappa])
        self.recog_kappa_mean, self.recog_kappa_prec_values = ParallelDense(n_x), tf.keras.Sequential([clip_by_value, ParallelTriangularDense(n_x)])

        self.recog = tf.keras.models.Sequential([ParallelDense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_recog])
        self.recog_mean, self.recog_prec_values = ParallelDense(n_z), tf.keras.Sequential([clip_by_value, ParallelTriangularDense(n_z)])

        self.gener_z = tf.keras.models.Sequential([ParallelDense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_gener_z])
        self.gener_z_mean, self.gener_z_prec_values = ParallelDense(n_z), tf.keras.Sequential([clip_by_value, ParallelTriangularDense(n_z)])

        self.gener_kappa = tf.keras.models.Sequential([ParallelDense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_gener_kappa])
        self.gener_kappa_mean, self.gener_kappa_prec_values = ParallelDense(n_kappa_max), tf.keras.Sequential([clip_by_value, ParallelDense(n_kappa_max)])

        self.gener_y = tf.keras.models.Sequential([tf.keras.layers.Dense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_gener_y])
        self.gener_y_mean, self.gener_y_prec_values = tf.keras.layers.Dense(n_y), tf.keras.Sequential([clip_by_value, tf.keras.layers.Dense(n_y)])

        self.log_lambda = tf.keras.models.Sequential([ParallelDense(n_layer, activation='leaky_relu', kernel_initializer='he_uniform') for n_layer in n_scale] + [ParallelDense(1)])

        self.recog_y.build((None, None, n_y)), self.recog_y_mean.build((None, None, n_recog_y[-1])), self.recog_y_prec_values.build((None, None, n_recog_y[-1]))
        self.recog_kappa.build((None, n_site, n_kappa_max)), self.recog_kappa_mean.build((None, n_site, n_recog_kappa[-1])), self.recog_kappa_prec_values.build((None, n_site, n_recog_kappa[-1]))

        self.recog.build((None, None, n_site, n_x + n_kappa_max)), self.recog_mean.build((None, None, n_site, n_recog[-1])), self.recog_prec_values.build((None, None, n_site, n_recog[-1]))

        self.gener_z.build((None, None, n_site, n_x)), self.gener_z_mean.build((None, None, n_site, n_gener_z[-1])), self.gener_z_prec_values.build((None, None, n_site, n_gener_z[-1]))
        self.gener_kappa.build((None, None, n_site, n_z)), self.gener_kappa_mean.build((None, None, n_site, n_gener_kappa[-1])), self.gener_kappa_prec_values.build((None, None, n_site, n_gener_kappa[-1]))

        self.gener_y.build((None, None, n_x)), self.gener_y_mean.build((None, None, n_gener_y[-1])), self.gener_y_prec_values.build((None, None, n_gener_y[-1]))

        self.log_lambda.build((None, n_site, n_x))

        recog_layers = [self.recog, self.recog_mean, self.recog_prec_values]
        gener_layers = [self.gener_z, self.gener_z_mean, self.gener_z_prec_values, self.gener_kappa, self.gener_kappa_mean, self.gener_kappa_prec_values]

        recog_y_layers = [self.recog_y, self.recog_y_mean, self.recog_y_prec_values]
        recog_kappa_layers = [self.recog_kappa, self.recog_kappa_mean, self.recog_kappa_prec_values]

        gener_y_layers = [self.gener_y, self.gener_y_mean, self.gener_y_prec_values]

        self.lambda_recog_variables = sum([layer.variables for layer in recog_layers], [])
        self.lambda_gener_variables = sum([layer.variables for layer in gener_layers], [])
        self.lambda_scale_variables = self.log_lambda.variables

        self.eta_variables_3 = sum([layer.variables for layer in gener_y_layers], [])
        with tf.device('/cpu:0'):
            self.eta_variables_2 = [self.recog_x_init_mean, self.recog_x_init_prec_values, self.recog_x_trans_mat_values, self.gener_x_init_mean, self.gener_x_init_prec_values, self.gener_x_trans_mat_values, self.gener_x_trans_mean, self.gener_x_trans_prec_values]
        self.eta_variables_1 = sum([layer.variables for layer in recog_y_layers + recog_kappa_layers], []) + [self.recog_kappa_g_mean, self.recog_kappa_g_prec_values]

        self.lambda_optimizer = tf.keras.optimizers.legacy.Adam(self.learning_rate)
        self.eta_optimizer_3 = tf.keras.optimizers.legacy.Adam(self.learning_rate)
        with tf.device('/cpu:0'):
            self.eta_optimizer_2 = tf.keras.optimizers.legacy.Adam(self.learning_rate)
        self.eta_optimizer_1 = tf.keras.optimizers.legacy.Adam(self.learning_rate)

        self._train, self._eval_log_likelihood, self._eval_log_likelihood_x, self._decode_init, self._decode_update = _train, _eval_log_likelihood, _eval_log_likelihood_x, _decode_init, _decode_update

    def train(self, kappas_padded, masks, ys_padded, n_monte):
        return self._train(self, kappas_padded, masks, ys_padded, n_monte)


    def eval_log_likelihood(self, kappas, masks, y, log_scale_train, log_scale_eval, n_monte):
        return self._eval_log_likelihood(self, kappas, masks, y, log_scale_train, log_scale_eval, n_monte)


    def eval_log_likelihood_x(self, kappas, masks, y, log_scale_train, log_scale_eval, n_monte):
        return self._eval_log_likelihood_x(self, kappas, masks, y, log_scale_train, log_scale_eval, n_monte)


    def decode_init(self, kappas, masks, log_scale_train, log_scale_eval, n_monte, n_iter=0.):
        return self._decode_init(self, kappas, masks, log_scale_train, log_scale_eval, n_monte, n_iter)


    def decode_update(self, x_current_means, x_current_covs, kappas, masks, log_scale_train, log_scale_eval, n_monte):
        return self._decode_update(self, x_current_means, x_current_covs, kappas, masks, log_scale_train, log_scale_eval, n_monte)


    def decode(self, kappas, masks, log_scale_train, log_scale_eval, n_monte, n_iter_init, n_iter_update):

        y_smoothed_means, y_smoothed_covs, x_smoothed_means, x_smoothed_covs = self.decode_init(kappas, masks, log_scale_train, log_scale_eval, n_monte, n_iter_init)
        for i in range(n_iter_update):
            y_smoothed_means, y_smoothed_covs, x_smoothed_means, x_smoothed_covs = self.decode_update(x_smoothed_means, x_smoothed_covs, kappas, masks, log_scale_train, log_scale_eval, n_monte)

        return y_smoothed_means, y_smoothed_covs, x_smoothed_means, x_smoothed_covs



def _train(model, batch_kappas, batch_masks, batch_ys, n_monte):

    n_batch, R = tf.shape(batch_ys)[0], tf.shape(batch_ys)[1]
    n_x = model.n_params['n_x']
    n_site = model.n_params['n_site']
    n_compile = model.n_params['n_compile']

    batch_ns = tf.math.count_nonzero(tf.reduce_sum(batch_masks, (2, 3)), axis=-1)
    batch_us = calculateUpperBound(batch_ns, n_compile)

    def _train_batch(state, elem):

        kappas, masks, u, y = elem

        epsx = tf.random.normal((R, 1, n_x), dtype=tf.float32)

        y_not_nan, is_not_nan, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde = compute_eta_loss_1(model, kappas, masks, y)
        with tf.device('/cpu:0'):
            _, x = compute_eta_loss_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx)

        eta_cost_3, x_grad_1 = _apply_eta_gradients_3(model, y_not_nan, is_not_nan, x)

        lambda_recog_cost, lambda_gener_cost, x_grad_2 = _apply_lambda_gradients(model, kappas[:u], masks[:u], x, n_monte)

        x_grad = x_grad_1 + x_grad_2

        with tf.device('/cpu:0'):
            eta_cost_2, x_affine_diag_grad, x_affine_tri_grad, recog_x_prec_tilde_grad, recog_x_mean_dot_prec_tilde_grad = _apply_eta_gradients_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx, x_grad)
        eta_cost_1 = _apply_eta_gradients_1(model, kappas, masks, y, x_affine_diag_grad, x_affine_tri_grad, recog_x_prec_tilde_grad, recog_x_mean_dot_prec_tilde_grad)

        cost = eta_cost_1 + eta_cost_2 + eta_cost_3 + lambda_recog_cost + lambda_gener_cost

        return cost

    costs = tf.scan(_train_batch, elems=(batch_kappas, batch_masks, batch_us, batch_ys), initializer=0.)
    avg_cost = tf.reduce_mean(costs)

    return avg_cost


@tf.function(jit_compile=True, reduce_retracing=True)
def _apply_lambda_gradients(model, kappas, masks, x, n_monte):

    alpha_1 = model.alpha_1

    R = tf.shape(x)[0]

    n_site = model.n_params['n_site']

    with tf.GradientTape(persistent=True) as tape:

        tape.watch(x)

        log_lambda = tf.squeeze(model.log_lambda(matmul_mask_tranpose_vec(masks, x)), -1)
        log_lambda_g = tf.squeeze(model.log_lambda(x), -1)

        lambda_recog_loss, lambda_gener_loss = _compute_lambda_loss(model, matmul_mask_tranpose_vec(masks, x), kappas, n_monte)

        lambda_recog_weights = (1 - alpha_1) * tf.stop_gradient(tf.exp(2. * tf.math.log_softmax((1 - alpha_1) * lambda_recog_loss, axis=1)))
        lambda_recog_weights += alpha_1 * tf.stop_gradient(tf.exp(tf.math.log_softmax((1 - alpha_1) * lambda_recog_loss, axis=1)))
        lambda_gener_weights = tf.stop_gradient(tf.exp(tf.math.log_softmax((1 - alpha_1) * lambda_gener_loss, axis=1)))

        lambda_recog_all = tf.reduce_sum(tf.math.multiply_no_nan(lambda_recog_loss, lambda_recog_weights), axis=1)
        lambda_gener_all = tf.reduce_sum(tf.math.multiply_no_nan(lambda_gener_loss, lambda_gener_weights), axis=1)

        lambda_recog_cost = - tf.reduce_sum(lambda_recog_all * tf.reduce_sum(masks, axis=1))
        lambda_gener_cost = - tf.reduce_sum(lambda_gener_all * tf.reduce_sum(masks, axis=1))
        lambda_scale_cost = - tf.reduce_sum(log_lambda * tf.reduce_sum(masks, axis=1)) + tf.reduce_sum(tf.exp(log_lambda_g))

    lambda_recog_gradients = tape.gradient(lambda_recog_cost, model.lambda_recog_variables)
    lambda_gener_gradients = tape.gradient(lambda_gener_cost, model.lambda_gener_variables)
    lambda_scale_gradients = tape.gradient(lambda_scale_cost, model.lambda_scale_variables)

    x_grad = tape.gradient(lambda_recog_cost, x) + tape.gradient(lambda_gener_cost, x) + tape.gradient(lambda_scale_cost, x)

    if not tf.math.reduce_any([tf.math.reduce_any([tf.math.is_nan(grad), tf.math.is_inf(grad)]) for grad in lambda_recog_gradients + lambda_gener_gradients + lambda_scale_gradients if grad is not None]):

        model.lambda_optimizer.apply_gradients(zip(lambda_recog_gradients, model.lambda_recog_variables))
        model.lambda_optimizer.apply_gradients(zip(lambda_gener_gradients, model.lambda_gener_variables))
        model.lambda_optimizer.apply_gradients(zip(lambda_scale_gradients, model.lambda_scale_variables))

    return lambda_recog_cost, lambda_gener_cost, x_grad


def _compute_lambda_loss(model, x, kappas, n_monte):

    n = tf.shape(x)[0]
    n_z, n_x = model.n_params['n_z'], model.n_params['n_x']
    n_site = model.n_params['n_site']

    n_kappa_max = tf.reduce_max(model.n_params['n_kappa'])

    gz_z = model.gener_z(x[:, tf.newaxis])
    zp_mean, zp_prec_values = model.gener_z_mean(gz_z), model.gener_z_prec_values(gz_z)
    zp_prec_diag, zp_prec_tri = zp_prec_values[:, :, :, :n_z], fill_triangular(zp_prec_values[:, :, :, n_z:])

    rz = model.recog(tf.concat([x, kappas], axis=-1)[:, tf.newaxis])
    z_mean, z_prec_values = model.recog_mean(rz), model.recog_prec_values(rz)
    z_prec_diag, z_prec_tri = z_prec_values[:, :, :, :n_z], fill_triangular(z_prec_values[:, :, :, n_z:])

    eps = tf.random.normal((n, n_monte, n_site, n_z))
    z = z_mean + tf.squeeze(tf.linalg.triangular_solve(tf.linalg.matrix_transpose(z_prec_tri), (eps * tf.exp(- 0.5 * z_prec_diag))[:, :, :, :, tf.newaxis], lower=False), axis=-1)

    #  recog

    gz_kappa = model.gener_kappa(z)
    kappa_mean, kappa_prec_diag = model.gener_kappa_mean(gz_kappa), model.gener_kappa_prec_values(gz_kappa)

    reconstr_kappa_loss = tf.reduce_sum(tf.math.multiply_no_nan(- 0.5 * tf.math.log(2 * np.pi) + 0.5 * kappa_prec_diag - 0.5 * tf.square((kappas[:, tf.newaxis] - kappa_mean) * tf.exp(0.5 * kappa_prec_diag)), model.n_kappa_masks), -1)
    latent_loss = 0.5 * tf.reduce_sum(tf.stop_gradient(zp_prec_diag), -1) - 0.5 * tf.reduce_sum(tf.square(tf.linalg.matvec(tf.linalg.matrix_transpose(tf.stop_gradient(zp_prec_tri)), z - tf.stop_gradient(zp_mean)) * tf.exp(0.5 * tf.stop_gradient(zp_prec_diag))), -1) - 0.5 * tf.reduce_sum(tf.stop_gradient(z_prec_diag), -1) + 0.5 * tf.reduce_sum(tf.square(tf.linalg.matvec(tf.linalg.matrix_transpose(tf.stop_gradient(z_prec_tri)), z - tf.stop_gradient(z_mean)) * tf.exp(0.5 * tf.stop_gradient(z_prec_diag))), -1)

    lambda_recog_loss = reconstr_kappa_loss + latent_loss

    # gener

    gz_kappa = model.gener_kappa(tf.stop_gradient(z))
    kappa_mean, kappa_prec_diag = model.gener_kappa_mean(gz_kappa), model.gener_kappa_prec_values(gz_kappa)

    reconstr_kappa_loss = tf.reduce_sum(tf.math.multiply_no_nan(- 0.5 * tf.math.log(2 * np.pi) + 0.5 * kappa_prec_diag - 0.5 * tf.square((kappas[:, tf.newaxis] - kappa_mean) * tf.exp(0.5 * kappa_prec_diag)), model.n_kappa_masks), -1)
    latent_loss = 0.5 * tf.reduce_sum(zp_prec_diag, -1) - 0.5 * tf.reduce_sum(tf.square(tf.linalg.matvec(tf.linalg.matrix_transpose(zp_prec_tri), tf.stop_gradient(z) - zp_mean) * tf.exp(0.5 * zp_prec_diag)), -1) - tf.stop_gradient(0.5 * tf.reduce_sum(tf.stop_gradient(z_prec_diag), -1) - 0.5 * tf.reduce_sum(tf.square(eps), -1))

    lambda_gener_loss = reconstr_kappa_loss + latent_loss

    return lambda_recog_loss, lambda_gener_loss


@tf.function(experimental_compile=True)
def _apply_eta_gradients_3(model, y_not_nan, is_not_nan, x):

    with tf.GradientTape(persistent=True) as tape:
        tape.watch(x)
        eta_cost_3 = _compute_eta_loss_3(model, y_not_nan, is_not_nan, x)

    x_grad = tape.gradient(eta_cost_3, x)

    eta_gradients_3 = tape.gradient(eta_cost_3, model.eta_variables_3)
    if not tf.math.reduce_any([tf.math.reduce_any([tf.math.is_nan(grad), tf.math.is_inf(grad)]) for grad in eta_gradients_3 if grad is not None]):
        model.eta_optimizer_3.apply_gradients(zip(eta_gradients_3, model.eta_variables_3))
    return eta_cost_3, x_grad


@tf.function(experimental_compile=True)
def _apply_eta_gradients_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx, x_grad):

    with tf.GradientTape(persistent=True) as tape:
        tape.watch([x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde])
        eta_cost_2, x = _compute_eta_loss_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx)
        eta_cost_2 += tf.reduce_sum(x_grad * x)

    x_affine_diag_grad, x_affine_tri_grad, recog_x_prec_tilde_grad, recog_x_mean_dot_prec_tilde_grad = tape.gradient(eta_cost_2, [x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde])

    eta_gradients_2 = tape.gradient(eta_cost_2, model.eta_variables_2)
    if not tf.math.reduce_any([tf.math.reduce_any([tf.math.is_nan(grad), tf.math.is_inf(grad)]) for grad in eta_gradients_2 if grad is not None]):
        model.eta_optimizer_2.apply_gradients(zip(eta_gradients_2, model.eta_variables_2))

    return eta_cost_2, x_affine_diag_grad, x_affine_tri_grad, recog_x_prec_tilde_grad, recog_x_mean_dot_prec_tilde_grad


@tf.function(experimental_compile=True)
def _apply_eta_gradients_1(model, kappas, masks, y, x_affine_diag_grad, x_affine_tri_grad, recog_x_prec_tilde_grad, recog_x_mean_dot_prec_tilde_grad):

    with tf.GradientTape(persistent=True) as tape:
        _, _, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde = _compute_eta_loss_1(model, kappas, masks, y)
        eta_cost_1 = tf.reduce_sum(x_affine_diag_grad * x_affine_diag) + tf.reduce_sum(x_affine_tri_grad * x_affine_tri) + tf.reduce_sum(recog_x_prec_tilde_grad * recog_x_prec_tilde) + tf.reduce_sum(recog_x_mean_dot_prec_tilde_grad * recog_x_mean_dot_prec_tilde)

    eta_gradients_1 = tape.gradient(eta_cost_1, model.eta_variables_1)
    if not tf.math.reduce_any([tf.math.reduce_any([tf.math.is_nan(grad), tf.math.is_inf(grad)]) for grad in eta_gradients_1 if grad is not None]):
        model.eta_optimizer_1.apply_gradients(zip(eta_gradients_1, model.eta_variables_1))

    return eta_cost_1


def _compute_eta_loss_1(model, kappas, masks, y):

    R = tf.shape(y)[0]
    n_x, n_y = model.n_params['n_x'], model.n_params['n_y']
    n_site = model.n_params['n_site']

    is_not_nan = tf.dtypes.cast(tf.math.logical_not(tf.reduce_any(tf.math.is_nan(y), -1, keepdims=True)), dtype=tf.float32)
    y_not_nan = tf.math.multiply_no_nan(y, tf.dtypes.cast(tf.math.logical_not(tf.math.is_nan(y)), dtype=tf.float32))

    x_affine_diag, x_affine_tri = model.x_affine_values[:n_x], fill_triangular(model.x_affine_values[n_x:])

    rx = model.recog_y(y_not_nan)
    recog_x_mean, recog_x_prec_values = model.recog_y_mean(rx), model.recog_y_prec_values(rx)
    recog_x_prec_diag, recog_x_prec_tri = recog_x_prec_values[:, :, :n_x], fill_triangular(recog_x_prec_values[:, 0, n_x:])

    recog_x_prec_tri_tilde = tf.exp(0.5 * x_affine_diag[:, tf.newaxis] + 0.5 * recog_x_prec_diag) * tf.linalg.matmul(x_affine_tri, recog_x_prec_tri)
    recog_x_prec_tilde = tf.matmul(recog_x_prec_tri_tilde, tf.linalg.matrix_transpose(recog_x_prec_tri_tilde))
    recog_x_mean_dot_prec_tilde = tf.matmul(tf.matmul(recog_x_mean, recog_x_prec_tri) * tf.exp(0.5 * recog_x_prec_diag), tf.linalg.matrix_transpose(recog_x_prec_tri_tilde))

    recog_x_prec_tilde, recog_x_mean_dot_prec_tilde = tf.math.multiply_no_nan(recog_x_prec_tilde, is_not_nan), tf.math.multiply_no_nan(recog_x_mean_dot_prec_tilde, is_not_nan)

    rlambda = model.recog_kappa(kappas)
    lambda_mean, lambda_prec_values = model.recog_kappa_mean(rlambda), model.recog_kappa_prec_values(rlambda)
    lambda_prec_diag, lambda_prec_tri = lambda_prec_values[:, :, :n_x], fill_triangular(lambda_prec_values[:, :, n_x:])

    lambda_prec_tri_tilde = tf.exp(0.5 * x_affine_diag[:, tf.newaxis] + 0.5 * lambda_prec_diag[:, :, tf.newaxis]) * tf.matmul(x_affine_tri, lambda_prec_tri)
    lambda_prec_tilde = tf.matmul(lambda_prec_tri_tilde, tf.linalg.matrix_transpose(lambda_prec_tri_tilde))
    lambda_mean_dot_prec_tilde = tf.linalg.matvec(lambda_prec_tri_tilde, tf.linalg.matvec(tf.linalg.matrix_transpose(lambda_prec_tri), lambda_mean) * tf.exp(0.5 * lambda_prec_diag))

    lambda_g_mean = tf.reshape(model.recog_kappa_g_mean, [1, n_x])
    lambda_g_prec_diag, lambda_g_prec_tri = model.recog_kappa_g_prec_values[:n_x], fill_triangular(model.recog_kappa_g_prec_values[n_x:])

    lambda_g_prec_tri_tilde = tf.exp(0.5 * x_affine_diag[:, tf.newaxis] + 0.5 * lambda_g_prec_diag) * tf.linalg.matmul(x_affine_tri, lambda_g_prec_tri)
    lambda_g_prec_tilde = tf.matmul(lambda_g_prec_tri_tilde, tf.linalg.matrix_transpose(lambda_g_prec_tri_tilde))
    lambda_g_mean_dot_prec_tilde = tf.matmul(tf.matmul(lambda_g_mean, lambda_g_prec_tri) * tf.exp(0.5 * lambda_g_prec_diag), tf.linalg.matrix_transpose(lambda_g_prec_tri_tilde))

    recog_x_prec_tilde += tf.reduce_sum(matmul_mask_mat(masks, lambda_prec_tilde), axis=1) + lambda_g_prec_tilde
    recog_x_mean_dot_prec_tilde += tf.reduce_sum(matmul_mask_vec(masks, lambda_mean_dot_prec_tilde), axis=1, keepdims=True) + lambda_g_mean_dot_prec_tilde

    return y_not_nan, is_not_nan, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde



def _compute_eta_loss_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx):

    n_x, n_y = model.n_params['n_x'], model.n_params['n_y']
    n_site = model.n_params['n_site']

    R = tf.shape(recog_x_prec_tilde)[0]

    recog_x_trans_mat = tf.reshape(model.recog_x_trans_mat_values, [n_x, n_x])
    recog_x_trans_mean = tf.reshape(model.recog_x_trans_mean, (1, n_x))
    recog_x_trans_prec = tf.eye(n_x)

    recog_x_init_prec_diag, recog_x_init_prec_tri = model.recog_x_init_prec_values[:n_x], fill_triangular(model.recog_x_init_prec_values[n_x:])
    recog_x_init_prec = tf.matmul(recog_x_init_prec_tri * tf.exp(0.5 * recog_x_init_prec_diag), tf.transpose(recog_x_init_prec_tri * tf.exp(0.5 * recog_x_init_prec_diag)))
    recog_x_init_mean = tf.reshape(model.recog_x_init_mean, (1, n_x))

    recog_x_grads = tf.concat([recog_x_mean_dot_prec_tilde[:1] + tf.matmul(recog_x_init_mean, recog_x_init_prec) - tf.matmul(recog_x_trans_mean, tf.transpose(recog_x_trans_mat)), recog_x_mean_dot_prec_tilde[1:-1] + recog_x_trans_mean - tf.matmul(recog_x_trans_mean, tf.transpose(recog_x_trans_mat)), recog_x_mean_dot_prec_tilde[-1:] + recog_x_trans_mean], axis=0)

    x_cholesky_diags, x_cholesky_off_diags, vs_tilde, ws_tilde = _cholesky_update(recog_x_prec_tilde, recog_x_grads, recog_x_trans_mat, recog_x_trans_prec, recog_x_init_prec, epsx)
    x_tilde = vs_tilde + ws_tilde

    gener_x_trans_mat = tf.reshape(model.gener_x_trans_mat_values, [n_x, n_x])
    gener_x_trans_mean = tf.reshape(model.gener_x_trans_mean, (1, n_x))
    gener_x_trans_prec_diag, gener_x_trans_prec_tri = model.gener_x_trans_prec_values[:n_x], fill_triangular(model.gener_x_trans_prec_values[n_x:])

    gener_x_init_prec_diag, gener_x_init_prec_tri = model.gener_x_init_prec_values[:n_x], fill_triangular(model.gener_x_init_prec_values[n_x:])
    gener_x_init_mean = tf.reshape(model.gener_x_init_mean, (1, n_x))

    latent_loss = - tf.reduce_sum(tf.math.log(tf.linalg.diag_part(x_cholesky_diags)), -1, keepdims=True) + 0.5 * tf.reduce_sum(tf.square(epsx), -1)
    latent_loss += tf.concat([0.5 * tf.reduce_sum(gener_x_init_prec_diag) - 0.5 * tf.reduce_sum(tf.square(tf.matmul(x_tilde[:1] - gener_x_init_mean, gener_x_init_prec_tri) * tf.exp(0.5 * gener_x_init_prec_diag)), -1), 0.5 * tf.reduce_sum(gener_x_trans_prec_diag) - 0.5 * tf.reduce_sum(tf.square(tf.matmul(x_tilde[1:] - tf.matmul(x_tilde[:-1], gener_x_trans_mat) - gener_x_trans_mean, gener_x_trans_prec_tri) * tf.exp(0.5 * gener_x_trans_prec_diag)), -1)], axis=0)

    eta_cost_2 = - tf.reduce_sum(latent_loss)

    x = tf.linalg.matrix_transpose(tf.matmul(tf.transpose(x_affine_tri), tf.linalg.matrix_transpose(x_tilde * tf.exp(0.5 * x_affine_diag))))

    return eta_cost_2, x



def _compute_eta_loss_3(model, y_not_nan, is_not_nan, x):

    n_x, n_y = model.n_params['n_x'], model.n_params['n_y']
    n_site = model.n_params['n_site']

    gy = model.gener_y(x)
    y_mean, y_prec_diag = model.gener_y_mean(gy), model.gener_y_prec_values(gy)

    reconstr_y_loss = tf.reduce_sum(- 0.5 * tf.math.log(2 * np.pi) + 0.5 * y_prec_diag - 0.5 * tf.square((y_not_nan - y_mean) * tf.exp(0.5 * y_prec_diag)) , -1)
    reconstr_y_loss = tf.math.multiply_no_nan(reconstr_y_loss, tf.squeeze(is_not_nan, -1))

    eta_cost_3 = - tf.reduce_sum(reconstr_y_loss)

    return eta_cost_3


compute_eta_loss_1 = tf.function(_compute_eta_loss_1, jit_compile=True)
compute_eta_loss_2 = tf.function(_compute_eta_loss_2, jit_compile=True)
compute_eta_loss_3 = tf.function(_compute_eta_loss_3, jit_compile=True)

compute_lambda_loss = tf.function(_compute_lambda_loss, jit_compile=True)


def compute_lambda_loss_split(model, x, kappas, n_monte, split_num=1):

    split_size = tf.shape(kappas)[0] // split_num + 1

    lambda_loss = []

    for i in range(split_num):
        x_splited, kappas_splited = x[i*split_size:(i+1)*split_size], kappas[i*split_size:(i+1)*split_size]
        lambda_loss_splited, _ = compute_lambda_loss(model, x_splited, kappas_splited, n_monte)
        lambda_loss.append(lambda_loss_splited)

    lambda_loss = tf.concat(lambda_loss, axis=0)

    return lambda_loss


@tf.function
def _eval_log_likelihood(model, kappas, masks, y, log_scale_train, log_scale_eval, n_monte):

    R = tf.shape(y)[0]
    n_x = model.n_params['n_x']
    n_site = model.n_params['n_site']

    alpha_1 = model.alpha_1

    epsx = tf.random.normal((R, 1, n_x), dtype=tf.float32)

    y_not_nan, is_not_nan, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde = compute_eta_loss_1(model, kappas, masks, y)
    with tf.device('/cpu:0'):
        eta_cost_2, x = compute_eta_loss_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx)
    eta_cost_3 = compute_eta_loss_3(model, y_not_nan, is_not_nan, x)

    log_likelihood = - eta_cost_2 - eta_cost_3

    log_lambda = tf.squeeze(model.log_lambda(matmul_mask_tranpose_vec(masks, x)), axis=-1)
    log_lambda_g = tf.squeeze(model.log_lambda(x), axis=-1)

    lambda_loss = compute_lambda_loss_split(model, matmul_mask_tranpose_vec(masks, x), kappas, n_monte, n_site)

    lambda_scale_all = log_scale_train + log_lambda + (1 - alpha_1) ** -1 * tfp.math.reduce_logmeanexp((1 - alpha_1) * lambda_loss, axis=1)
    lambda_g_scale_all = tf.exp(log_scale_train + log_lambda_g - log_scale_eval)

    log_likelihood += tf.reduce_sum(tf.math.multiply_no_nan(tf.reduce_sum(masks, axis=1), lambda_scale_all)) - tf.reduce_sum(lambda_g_scale_all)

    return log_likelihood


@tf.function
def _eval_log_likelihood_x(model, kappas, masks, y, log_scale_train, log_scale_eval, n_monte):

    R = tf.shape(y)[0]
    n_x = model.n_params['n_x']
    n_site = model.n_params['n_site']

    alpha_1 = model.alpha_1

    epsx = tf.random.normal((R, 1, n_x), dtype=tf.float32)

    y_not_nan, is_not_nan, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde = compute_eta_loss_1(model, kappas, masks, y)
    with tf.device('/cpu:0'):
        eta_cost_2, x = compute_eta_loss_2(model, x_affine_diag, x_affine_tri, recog_x_prec_tilde, recog_x_mean_dot_prec_tilde, epsx)
    eta_cost_3 = compute_eta_loss_3(model, y_not_nan, is_not_nan, x)

    log_likelihood_x = - eta_cost_2 - eta_cost_3

    log_lambda = tf.squeeze(model.log_lambda(matmul_mask_tranpose_vec(masks, x)), axis=-1)
    log_lambda_g = tf.squeeze(model.log_lambda(x), axis=-1)

    lambda_scale_all = log_scale_train + log_lambda
    lambda_g_scale_all = tf.exp(log_scale_train + log_lambda_g - log_scale_eval)

    log_likelihood_x += tf.reduce_sum(tf.math.multiply_no_nan(tf.reduce_sum(masks, axis=1), lambda_scale_all)) - tf.reduce_sum(lambda_g_scale_all)

    return log_likelihood_x



def _eval_log_lambda_derivative(model, x, kappas, log_scale_train, n_monte):

    alpha_1 = model.alpha_1

    n_x = model.n_params['n_x']

    x_flattened = tf.reshape(x, (-1, n_x))
    n_site = model.n_params['n_site']

    with tf.GradientTape() as tape2:
        tape2.watch(x_flattened)
        with tf.GradientTape(persistent=True) as tape1:
            tape1.watch(x_flattened)

            x = tf.reshape(x_flattened, (-1, n_site, n_x))

            log_lambda = tf.squeeze(model.log_lambda(x), axis=-1)
            lambda_recog_loss, lambda_gener_loss, = _compute_lambda_loss(model, x, kappas, n_monte)

            loss = log_scale_train + log_lambda[:, tf.newaxis] + lambda_gener_loss
            weights = tf.stop_gradient(tf.math.softmax((1 - alpha_1) * loss, axis=1))

            cost1 = tf.reduce_sum(weights * loss)
            cost2 = tf.reduce_sum(weights * (tf.exp(loss - tf.stop_gradient(loss)) - loss))

        grad1 = tape1.gradient(cost1, x_flattened)
        grad2 = tape1.gradient(cost2, x_flattened)

    hessian = tape2.batch_jacobian(grad2, x_flattened)

    grad_log_lambda_x = tf.reshape(grad1, (-1, n_site, n_x))

    hessian_log_lambda_x = - (1 - alpha_1) * grad_log_lambda_x[:, :, :, tf.newaxis] * grad_log_lambda_x[:, :, tf.newaxis] - alpha_1 * tf.reshape(hessian, (-1, n_site, n_x, n_x))

    return grad_log_lambda_x, hessian_log_lambda_x


def _eval_lambda_g_derivative(model, x, log_scale_train, n_monte):

    n_x = model.n_params['n_x']
    n_site = model.n_params['n_site']

    with tf.GradientTape() as tape:
        tape.watch(x)

        log_lambda_g = tf.squeeze(model.log_lambda(x), axis=-1)
        loss = log_scale_train + log_lambda_g

    grad_log_lambda_x = tape.gradient(loss, x)

    grad_lambda_x = tf.exp(loss)[:, :, tf.newaxis] * grad_log_lambda_x
    hessian_lambda_x = tf.exp(loss)[:, :, tf.newaxis, tf.newaxis] * grad_log_lambda_x[:, :, tf.newaxis] * grad_log_lambda_x[:, :, :, tf.newaxis]

    return grad_lambda_x, hessian_lambda_x


@tf.function(experimental_compile=True)
def eval_log_lambda_derivative(model, x, kappas, log_scale_train, n_monte):
    return _eval_log_lambda_derivative(model, x, kappas, log_scale_train, n_monte)


@tf.function(experimental_compile=True)
def eval_lambda_g_derivative(model, x, log_scale_train, n_monte):
    return _eval_lambda_g_derivative(model, x, log_scale_train, n_monte)


def eval_log_lambda_derivative_split(model, x, kappas, log_scale_train, n_monte, split_num=1):

    split_size = tf.shape(kappas)[0] // split_num + 1

    grad_log_lambda_x, hessian_log_lambda_x = [], []

    for i in range(split_num):
        x_splited, kappas_splited = x[i*split_size:(i+1)*split_size], kappas[i*split_size:(i+1)*split_size]
        grad_log_lambda_x_splited, hessian_log_lambda_x_splited = eval_log_lambda_derivative(model, x_splited, kappas_splited, log_scale_train, n_monte)
        grad_log_lambda_x.append(grad_log_lambda_x_splited), hessian_log_lambda_x.append(hessian_log_lambda_x_splited)

    grad_log_lambda_x, hessian_log_lambda_x = tf.concat(grad_log_lambda_x, axis=0), tf.concat(hessian_log_lambda_x, axis=0)

    return grad_log_lambda_x, hessian_log_lambda_x


def eval_lambda_g_derivative_split(model, x, log_scale_train, n_monte, split_num=1):

    split_size = tf.shape(x)[0] // split_num + 1

    grad_lambda_x, hessian_lambda_x = [], []

    for i in range(split_num):
        x_splited = x[i*split_size:(i+1)*split_size]
        grad_lambda_x_splited, hessian_lambda_x_splited = eval_lambda_g_derivative(model, x_splited, log_scale_train, n_monte)
        grad_lambda_x.append(grad_lambda_x_splited), hessian_lambda_x.append(hessian_lambda_x_splited)

    grad_lambda_x, hessian_lambda_x = tf.concat(grad_lambda_x, axis=0), tf.concat(hessian_lambda_x, axis=0)

    return grad_lambda_x, hessian_lambda_x


@tf.function
def _decode_init(model, kappas, masks, log_scale_train, log_scale_eval, n_monte, n_iter):

    n_site = model.n_params['n_site']
    n_kappa_max = tf.reduce_max(model.n_params['n_kappa'])
    R = tf.shape(masks)[1]

    inds = [tf.boolean_mask(tf.repeat(tf.range(R)[tf.newaxis], tf.shape(masks[:, :, site])[0], axis=0), masks[:, :, site]) for site in range(n_site)]
    n_max = tf.reduce_max([tf.reduce_max(tf.unique_with_counts(inds[site])[-1]) for site in range(n_site)])

    kappas_ragged = tf.stack([tf.RaggedTensor.from_value_rowids(tf.boolean_mask(tf.squeeze(kappas[:, site]), tf.reduce_sum(masks[:, :, site], axis=1)), inds[site], R).to_tensor(shape=(R, n_max, n_kappa_max)) for site in range(n_site)], axis=2)
    masks_ragged = tf.stack([tf.RaggedTensor.from_value_rowids(tf.ones_like(inds[site], dtype=tf.float32), inds[site], R).to_tensor(shape=(R, n_max)) for site in range(n_site)], axis=2)[:, :, tf.newaxis]

    n_x, n_y = model.n_params['n_x'], model.n_params['n_y']

    kappas_ragged = tf.repeat(kappas_ragged, n_iter, axis=0)
    masks_ragged = tf.repeat(masks_ragged, n_iter, axis=0)
    updates = tf.reshape(tf.concat([tf.zeros((R, n_iter-1)), tf.ones((R, 1))], axis=1), (R*n_iter, ))

    x_filtered_state = collections.namedtuple('x_filtered_state', ['x_filtered_mean', 'x_filtered_prec', 'x_predicted_mean', 'x_predicted_prec'])
    x_smoothed_state = collections.namedtuple('x_smoothed_state', ['x_smoothed_mean', 'x_smoothed_cov'])

    x_affine_diag, x_affine_tri = model.x_affine_values[:n_x], fill_triangular(model.x_affine_values[n_x:])

    x_init_prec_diag, x_init_prec_tri = model.gener_x_init_prec_values[:n_x], fill_triangular(model.gener_x_init_prec_values[n_x:])
    x_init_prec = tf.matmul(x_init_prec_tri * tf.exp(0.5 * x_init_prec_diag), tf.transpose(x_init_prec_tri * tf.exp(0.5 * x_init_prec_diag)))
    x_init_mean = tf.reshape(model.gener_x_init_mean, (1, n_x))

    x_trans_mean = tf.reshape(model.gener_x_trans_mean, (1, n_x))

    x_trans_mat = tf.reshape(model.gener_x_trans_mat_values, [n_x, n_x])
    x_trans_mat /= tf.nn.relu(tf.reduce_max(tf.abs(tf.linalg.eigvals(x_trans_mat))) - 1.) + 1.

    x_trans_prec_diag, x_trans_prec_tri = model.gener_x_trans_prec_values[:n_x], fill_triangular(model.gener_x_trans_prec_values[n_x:])
    x_trans_prec = tf.matmul(x_trans_prec_tri * tf.exp(0.5 * x_trans_prec_diag), tf.transpose(x_trans_prec_tri * tf.exp(0.5 * x_trans_prec_diag)))

    def _filter_predict(model, x_filtered_mean, x_filtered_prec):
        x_predicted_mean, x_predicted_prec = x_trans_mean + tf.matmul(x_filtered_mean, x_trans_mat), tf.linalg.inv(tf.matmul(tf.transpose(x_trans_mat), tf.linalg.solve(x_filtered_prec, x_trans_mat)) + tf.linalg.inv(x_trans_prec))
        return x_predicted_mean, x_predicted_prec

    def _filter_correct(model, kappa, mask, x_current_mean, x_predicted_mean, x_predicted_prec, log_scale_train, log_scale_eval, n_monte):

        x_current_mean_orig = tf.transpose(tf.matmul(tf.transpose(x_affine_tri), tf.transpose(x_current_mean * tf.exp(0.5 * x_affine_diag))))

        grad_log_lambda_x, hessian_log_lambda_x = eval_log_lambda_derivative(model, matmul_mask_tranpose_vec(mask, x_current_mean_orig[tf.newaxis]), kappa, log_scale_train, n_monte)
        grad_lambda_x, hessian_lambda_x = eval_lambda_g_derivative(model, tf.repeat(x_current_mean_orig[tf.newaxis], n_site, axis=1), log_scale_train, n_monte)

        likelihood_grads = tf.squeeze(tf.reduce_sum(- matmul_mask_vec(mask, grad_log_lambda_x) + grad_lambda_x * tf.exp(- log_scale_eval), axis=1, keepdims=True), axis=0)
        likelihood_hessians = tf.squeeze(tf.reduce_sum(- matmul_mask_mat(mask, hessian_log_lambda_x) + hessian_lambda_x * tf.exp(- log_scale_eval), axis=1), axis=0)

        likelihood_grads_tilde = tf.linalg.matrix_transpose(tf.matmul(x_affine_tri, tf.linalg.matrix_transpose(likelihood_grads))) * tf.exp(0.5 * x_affine_diag)
        likelihood_hessians_tilde = tf.exp(0.5 * x_affine_diag[:, tf.newaxis] + 0.5 * x_affine_diag) * tf.matmul(x_affine_tri, tf.linalg.matrix_transpose(tf.matmul(x_affine_tri, likelihood_hessians)))

        x_filtered_prec = x_predicted_prec + likelihood_hessians_tilde
        x_filtered_mean = x_current_mean + tf.transpose(tf.linalg.solve(x_predicted_prec + likelihood_hessians_tilde, tf.transpose(- likelihood_grads_tilde - tf.linalg.matmul(x_current_mean - x_predicted_mean, x_predicted_prec))))

        return x_filtered_mean, x_filtered_prec


    def _smooth_update(model, x_predicted_mean, x_predicted_prec, x_filtered_mean, x_filtered_prec, x_next_smoothed_mean, x_next_smoothed_cov):

        n_x = model.n_params['n_x']

        x_kalman_gain = tf.matmul(tf.linalg.solve(x_filtered_prec, x_trans_mat), x_predicted_prec)
        x_smoothed_mean = x_filtered_mean + tf.matmul(x_next_smoothed_mean - x_predicted_mean, tf.transpose(x_kalman_gain))
        x_smoothed_cov = tf.linalg.inv(x_filtered_prec + tf.matmul(tf.matmul(x_trans_mat, x_trans_prec), tf.transpose(x_trans_mat))) + tf.matmul(tf.matmul(x_kalman_gain, x_next_smoothed_cov), tf.transpose(x_kalman_gain))

        return x_smoothed_mean, x_smoothed_cov


    def update_forward_fn(state, elem):
        kappa, mask, update = elem

        x_filtered_mean, x_filtered_prec = _filter_correct(model, kappa, mask, state.x_filtered_mean, state.x_predicted_mean, state.x_predicted_prec, log_scale_train, log_scale_eval, n_monte)
        x_predicted_mean, x_predicted_prec = _filter_predict(model, x_filtered_mean, x_filtered_prec)

        return x_filtered_state(x_filtered_mean, x_filtered_prec, x_predicted_mean * update + state.x_predicted_mean * (1 - update), x_predicted_prec * update + state.x_predicted_prec * (1 - update))


    def update_backward_fn(state, elem):

        x_filtered_mean, x_filtered_prec, x_predicted_mean, x_predicted_prec = elem
        x_smoothed_mean, x_smoothed_cov = _smooth_update(model, x_predicted_mean, x_predicted_prec, x_filtered_mean, x_filtered_prec, state.x_smoothed_mean, state.x_smoothed_cov)

        return x_smoothed_state(x_smoothed_mean, x_smoothed_cov)


    x_filtered_means, x_filtered_precs, x_predicted_means, x_predicted_precs = tf.scan(update_forward_fn, elems=(kappas_ragged, masks_ragged, updates), initializer=x_filtered_state(x_init_mean, x_init_prec, x_init_mean, x_init_prec))
    x_filtered_means, x_filtered_precs, x_predicted_means, x_predicted_precs = x_filtered_means[n_iter-1::n_iter], x_filtered_precs[n_iter-1::n_iter], x_predicted_means[n_iter-1::n_iter], x_predicted_precs[n_iter-1::n_iter]

    x_smoothed_mean, x_smoothed_cov = x_filtered_means[-1], tf.linalg.inv(x_filtered_precs[-1])

    x_smoothed_means, x_smoothed_covs = tf.scan(update_backward_fn, elems=(x_filtered_means[:-1], x_filtered_precs[:-1], x_predicted_means[:-1], x_predicted_precs[:-1]), initializer=x_smoothed_state(x_smoothed_mean, x_smoothed_cov), reverse=True)
    x_smoothed_means, x_smoothed_covs = tf.concat([x_smoothed_means, x_smoothed_mean[tf.newaxis]], axis=0), tf.concat([x_smoothed_covs, x_smoothed_cov[tf.newaxis]], axis=0)
    x_smoothed_covs_cholesky = tf.linalg.cholesky(x_smoothed_covs)


    x_smoothed_means = tf.linalg.matrix_transpose(tf.matmul(tf.transpose(x_affine_tri), tf.linalg.matrix_transpose(x_smoothed_means * tf.exp(0.5 * x_affine_diag))))
    x_smoothed_covs_cholesky = tf.matmul(tf.transpose(x_affine_tri), x_smoothed_covs_cholesky * tf.exp(0.5 * x_affine_diag[:, tf.newaxis]))
    x_smoothed_covs = tf.matmul(x_smoothed_covs_cholesky, tf.linalg.matrix_transpose(x_smoothed_covs_cholesky))

    epsx = tf.random.normal((R, n_monte, n_x))
    x = x_smoothed_means + tf.matmul(epsx, tf.linalg.matrix_transpose(x_smoothed_covs_cholesky))

    gy = model.gener_y(x)
    y_mean, y_prec_diag = model.gener_y_mean(gy), model.gener_y_prec_values(gy)

    y_smoothed_means = tf.reduce_mean(y_mean, 1, keepdims=True)
    y_smoothed_covs = tf.reduce_mean(tf.matmul((y_smoothed_means - y_mean)[:, :, :, tf.newaxis], (y_smoothed_means - y_mean)[:, :, tf.newaxis]), 1) + tf.reduce_mean(tf.linalg.diag(tf.exp(- y_prec_diag)), axis=1)

    return y_smoothed_means, y_smoothed_covs, x_smoothed_means, x_smoothed_covs


@tf.function
def _decode_update(model, x_current_means, x_current_covs, kappas, masks, log_scale_train, log_scale_eval, n_monte):

    R = tf.shape(x_current_means)[0]

    n_site = model.n_params['n_site']
    n_x = model.n_params['n_x']
    n_y = model.n_params['n_y']

    x_affine_diag, x_affine_tri = model.x_affine_values[:n_x], fill_triangular(model.x_affine_values[n_x:])

    x_init_prec_diag, x_init_prec_tri = model.gener_x_init_prec_values[:n_x], fill_triangular(model.gener_x_init_prec_values[n_x:])
    x_init_prec = tf.matmul(x_init_prec_tri * tf.exp(0.5 * x_init_prec_diag), tf.transpose(x_init_prec_tri * tf.exp(0.5 * x_init_prec_diag)))
    x_init_mean = tf.reshape(model.gener_x_init_mean, (1, n_x))

    x_trans_prec_diag, x_trans_prec_tri = model.gener_x_trans_prec_values[:n_x], fill_triangular(model.gener_x_trans_prec_values[n_x:])
    x_trans_prec = tf.matmul(x_trans_prec_tri * tf.exp(0.5 * x_trans_prec_diag), tf.transpose(x_trans_prec_tri * tf.exp(0.5 * x_trans_prec_diag)))
    x_trans_mean = tf.reshape(model.gener_x_trans_mean, (1, n_x))

    x_trans_mat = tf.reshape(model.gener_x_trans_mat_values, [n_x, n_x])
    x_trans_mat /= tf.nn.relu(tf.reduce_max(tf.abs(tf.linalg.eigvals(x_trans_mat))) - 1.) + 1.

    grad_log_lambda_x, hessian_log_lambda_x = eval_log_lambda_derivative_split(model, matmul_mask_tranpose_vec(masks, x_current_means), kappas, log_scale_train, n_monte, n_site)
    grad_lambda_x, hessian_lambda_x = eval_lambda_g_derivative_split(model, tf.repeat(x_current_means, n_site, axis=1), log_scale_train, n_monte, n_site)

    likelihood_grads = tf.reduce_sum(- matmul_mask_vec(masks, grad_log_lambda_x) + grad_lambda_x * tf.exp(- log_scale_eval), axis=1, keepdims=True)
    likelihood_hessians = tf.reduce_sum(- matmul_mask_mat(masks, hessian_log_lambda_x) + hessian_lambda_x * tf.exp(- log_scale_eval), axis=1)

    likelihood_grads_tilde = tf.linalg.matrix_transpose(tf.matmul(x_affine_tri, tf.linalg.matrix_transpose(likelihood_grads))) * tf.exp(0.5 * x_affine_diag)
    likelihood_hessians_tilde = tf.exp(0.5 * x_affine_diag[:, tf.newaxis] + 0.5 * x_affine_diag) * tf.linalg.matmul(x_affine_tri, tf.linalg.matrix_transpose(tf.matmul(x_affine_tri, likelihood_hessians)))

    x_current_means_tilde = tf.linalg.matrix_transpose(tf.linalg.triangular_solve(tf.transpose(x_affine_tri), tf.linalg.matrix_transpose(x_current_means), lower=False)) * tf.exp(- 0.5 * x_affine_diag)
    y_tilde = tf.concat([tf.matmul(x_current_means_tilde[:1] - x_init_mean, x_init_prec), tf.matmul(x_current_means_tilde[1:] - tf.matmul(x_current_means_tilde[:-1], x_trans_mat) - x_trans_mean, x_trans_prec)], axis=0)
    gener_x_grads = - likelihood_grads_tilde - tf.concat([y_tilde[:-1] - tf.matmul(y_tilde[1:], tf.transpose(x_trans_mat)), y_tilde[-1:]], axis=0)

    epsx = tf.random.normal((R, n_monte, n_x))

    with tf.device('/cpu:0'):
        x_cholesky_diags, x_cholesky_off_diags, vs_tilde, ws_tilde = cholesky_update(likelihood_hessians_tilde, gener_x_grads, x_trans_mat, x_trans_prec, x_init_prec, epsx)
        x_next_covs_tilde, x_next_pair_covs_tilde = covariance_update(x_cholesky_diags, x_cholesky_off_diags)

    vs = tf.linalg.matrix_transpose(tf.matmul(tf.transpose(x_affine_tri), tf.linalg.matrix_transpose(vs_tilde * tf.exp(0.5 * x_affine_diag))))
    ws = tf.linalg.matrix_transpose(tf.matmul(tf.transpose(x_affine_tri), tf.linalg.matrix_transpose(ws_tilde * tf.exp(0.5 * x_affine_diag))))

    x_next_means = x_current_means + vs
    x_next_covs = tf.linalg.matrix_transpose(tf.matmul(tf.transpose(x_affine_tri), tf.linalg.matrix_transpose(tf.matmul(tf.transpose(x_affine_tri), x_next_covs_tilde * tf.exp(0.5 * x_affine_diag[:, tf.newaxis] + 0.5 * x_affine_diag)))))

    x = x_next_means + ws

    gy = model.gener_y(x)
    y_mean, y_prec_diag = model.gener_y_mean(gy), model.gener_y_prec_values(gy)

    y_smoothed_means = tf.reduce_mean(y_mean, 1, keepdims=True)
    y_smoothed_covs = tf.reduce_mean(tf.matmul((y_smoothed_means - y_mean)[:, :, :, tf.newaxis], (y_smoothed_means - y_mean)[:, :, tf.newaxis]), 1) + tf.reduce_mean(tf.linalg.diag(tf.exp(- y_prec_diag)), axis=1)

    return y_smoothed_means, y_smoothed_covs, x_next_means, x_next_covs

def _cholesky_update(x_hessian_diags, x_grads, x_trans_mat, x_trans_prec, x_init_prec, epsx):

    R, n_monte, n_x = tf.shape(epsx)[0], tf.shape(epsx)[1], tf.shape(epsx)[2]

    x_forward_state = collections.namedtuple('x_forward_state', ['x_predicted_prec', 'x_cholesky_diag', 'x_cholesky_off_diag', 'u'])
    x_forward_init_state = x_forward_state(x_init_prec, tf.zeros((n_x, n_x)), tf.zeros((n_x, n_x)), tf.zeros((1, n_x)))

    def update_forward_fn(state, elem):

        x_hessian_diag, x_grad, m = elem
        x_predicted_prec, x_cholesky_off_diag, u = state.x_predicted_prec, state.x_cholesky_off_diag, state.u

        x_cholesky_diag = tf.linalg.cholesky(x_predicted_prec + x_hessian_diag + m * tf.matmul(tf.matmul(x_trans_mat, x_trans_prec), tf.transpose(x_trans_mat)))
        u = tf.transpose(tf.linalg.triangular_solve(x_cholesky_diag, tf.transpose(x_grad - tf.matmul(u, tf.transpose(x_cholesky_off_diag)))))
        x_cholesky_off_diag = tf.transpose(tf.linalg.triangular_solve(x_cholesky_diag, - tf.matmul(x_trans_mat, x_trans_prec)))

        x_predicted_prec = x_trans_prec - tf.matmul(x_cholesky_off_diag, tf.transpose(x_cholesky_off_diag))

        return x_forward_state(x_predicted_prec, x_cholesky_diag, x_cholesky_off_diag, u)

    ms = tf.concat([tf.ones(R-1), tf.zeros(1)], axis=0)
    _, x_cholesky_diags, x_cholesky_off_diags, us = tf.scan(update_forward_fn, elems=(x_hessian_diags, x_grads, ms), initializer=x_forward_init_state)

    x_backward_state = collections.namedtuple('x_backward_state', ['vw'])
    x_backward_init_state = x_backward_state(tf.zeros((n_monte+1, n_x)))

    gs = tf.concat([us, epsx], axis=1)

    def update_backward_fn(state, elem):

        x_cholesky_diag, x_cholesky_off_diag, g = elem
        vw = state.vw
        vw = tf.transpose(tf.linalg.triangular_solve(tf.transpose(x_cholesky_diag), tf.transpose(g - tf.matmul(vw, x_cholesky_off_diag)), lower=False))

        return x_backward_state(vw)

    vws,  = tf.scan(update_backward_fn, elems=(x_cholesky_diags, x_cholesky_off_diags, gs), initializer=x_backward_init_state, reverse=True)
    vs, ws = vws[:, :1], vws[:, 1:]

    return x_cholesky_diags, x_cholesky_off_diags, vs, ws



def _covariance_update(x_cholesky_diags, x_cholesky_off_diags):

    n_x = tf.shape(x_cholesky_diags)[-1]

    x_backward_state = collections.namedtuple('x_backward_state', ['x_smoothed_cov', 'x_smoothed_pair_cov'])
    x_backward_init_state = x_backward_state(tf.zeros((n_x, n_x)), tf.zeros((n_x, n_x)))

    x_cholesky_inv_diags = tf.linalg.triangular_solve(x_cholesky_diags, tf.eye(n_x) + tf.zeros_like(x_cholesky_diags))

    def update_backward_fn(state, elem):

        x_cholesky_inv_diag, x_cholesky_off_diag = elem
        x_smoothed_cov, x_smoothed_pair_cov = state.x_smoothed_cov, state.x_smoothed_pair_cov

        x_smoothed_pair_cov = - tf.matmul(tf.matmul(x_smoothed_cov, x_cholesky_off_diag), x_cholesky_inv_diag)
        x_smoothed_cov = tf.matmul(tf.transpose(x_cholesky_inv_diag), x_cholesky_inv_diag - tf.matmul(tf.transpose(x_cholesky_off_diag), x_smoothed_pair_cov))

        return x_backward_state(x_smoothed_cov, x_smoothed_pair_cov)

    x_smoothed_covs, x_smoothed_pair_covs = tf.scan(update_backward_fn, elems=(x_cholesky_inv_diags, x_cholesky_off_diags), initializer=x_backward_init_state, reverse=True)

    return x_smoothed_covs, x_smoothed_pair_covs


@tf.function(experimental_compile=True)
def cholesky_update(x_hessian_diags, x_grads, x_trans_mat, x_trans_prec, x_init_prec, epsx):
    return _cholesky_update(x_hessian_diags, x_grads, x_trans_mat, x_trans_prec, x_init_prec, epsx)


@tf.function(experimental_compile=True)
def covariance_update(x_cholesky_diags, x_cholesky_off_diags):
    return _covariance_update(x_cholesky_diags, x_cholesky_off_diags)