<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def sample_block_tridiag(H_diag, H_upper_diag):
 """ helper function for sampling block tridiag gaussians. this is only for speed comparison with the solve approach. """

    T, D, _ = H_diag.shape
    assert H_diag.ndim == 3 and H_diag.shape[2] == D
    assert H_upper_diag.shape == (T - 1, D, D)

    J_init = J_11 = J_22 = np.zeros((D, D))
    h_init = h_1 = h_2 = np.zeros((D,))

    J_21 = np.swapaxes(H_upper_diag, -1, -2)
    J_node = H_diag
    h_node = np.zeros((T,D))

    y = info_sample(J_init, h_init, 0,
                    J_11, J_21, J_22, h_1, h_2, np.zeros((T-1)),
                    J_node, h_node, np.zeros(T))
    return y