def categorical_kl_div(input_vb, ref_vb):
    """
    kl_div = \sum ref * (log(ref) - log(input))
    variables needed:
            input_vb: [batch_size x state_dim]
              ref_vb: [batch_size x state_dim]
    returns:
           kl_div_vb: [batch_size x 1]
    """
    return (ref_vb * (ref_vb.log() - input_vb.log())).sum(1, keepdim=True)