Unpack a lower triangular Cholesky matrix from a neural network output in Tensorflow

chol_unpack.py

import tensorflow as tf

def unpack_cholesky(q, ndims):
    # Build the lower-triangular Cholesky from the flat buffer (assumes q shape is [batchsize, cholsize])
    chol_diag = tf.nn.softplus(q[:,:ndims])
    chol_offdiag = q[:,ndims:]
    chol_rows = []
    chol_start = 0
    chol_end = 1
    for i in xrange(ndims):
        pieces = []
        if i > 0:
            pieces.append(chol_offdiag[:,chol_start:chol_end])
            chol_start = chol_end
            chol_end += i+1
        pieces.append(chol_diag[:,i:i+1])
        if i < (ndims-1):
            pieces.append(tf.zeros([tf.shape(chol_diag)[0], ndims-i-1]))
        print i, pieces
        chol_rows.append(tf.concat(1, pieces))
    return tf.reshape(tf.concat(1, chol_rows), [tf.shape(chol_diag)[0], ndims, ndims])