Define `chol_params` as a flattened Cholesky factor PyTree that consists
of diagonal elements followed by the off-diagonal elements in row-major
order for n = dim * (dim + 1) / 2 elements.

The diagonal (first dim elements) are passed through a softplus function
to ensure positivity, crucial to maintain a valid covariance matrix

This parameterization allows for unconstrained optimization while
ensuring the resulting covariance matrix Sigma = CC^T is symmetric and
positive definite.

The `chol_params` are then reshaped into a lower triangular matrix
`chol_factor` using `jnp.tril` and `jnp.diag` functions.

Fix testing bug, add docstrings, and change softmax to exponential when
converting `chol_params` to `chol_factor` in `_unflatten_cholesky`.

Refactor `_unflatten_cholesky()` function to take `dim` argument instead
of infering it (dynamically) from the `chol_params` input vector. This
avoids JIT compilation issues.

Also update docstrings.

Add assert statements that verify full-rank VI recovers the true,
full-rank covariance matrix.