function vaeLoss(inputs, outputs) {
  return tf.tidy(() => {
    const originalDim = inputs.shape[1];
    const decoderOutput = outputs[0];
    const zMean = outputs[1];
    const zLogVar = outputs[2];

    // First we compute a 'reconstruction loss' terms. The goal of minimizing
    // this term is to make the model outputs match the input data.
    const reconstructionLoss =
        tf.losses.meanSquaredError(inputs, decoderOutput).mul(originalDim);

    // binaryCrossEntropy can be used as an alternative loss function
    // const reconstructionLoss =
    //  tf.metrics.binaryCrossentropy(inputs, decoderOutput).mul(originalDim);

    // Next we compute the KL-divergence between zLogVar and zMean, minimizing
    // this term aims to make the distribution of latent variable more normally
    // distributed around the center of the latent space.
    let klLoss = zLogVar.add(1).sub(zMean.square()).sub(zLogVar.exp());
    klLoss = klLoss.sum(-1).mul(-0.5);

    return reconstructionLoss.add(klLoss).mean();
  });
}
The custom loss function for VAE.

@param {tf.tensor} inputs the encoder inputs a batched image tensor
@param {[tf.tensor]} outputs the vae outputs, [decoderOutput,
    ...encoderOutputs]
@param {number} vaeOpts.originalDim number of dimensions in the original data