How to use the numdifftools.directionaldiff function in numdifftools

lambda x: brsa._loglike_AR1_diagV_fitU(x, XTX, XTDX, XTFX, YTY_diag,
                                               YTDY_diag, YTFY_diag, XTY, XTDY,
                                               XTFY, X0TX0, X0TDX0, X0TFX0,
                                               XTX0, XTDX0, XTFX0, X0TY, X0TDY,
                                               X0TFY, np.log(snr) * 2, l_idx,
                                               n_C, n_T, n_V, n_run, n_X0,
                                               idx_param_fitU, n_C)[0],
        param0_fitU,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitU wrt Cholesky factor incorrect')

    # Test on a random direction
    vec = np.random.randn(np.size(param0_fitU))
    vec = vec / np.linalg.norm(vec)
    dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_diagV_fitU(x, XTX, XTDX, XTFX, YTY_diag,
                                               YTDY_diag, YTFY_diag, XTY, XTDY,
                                               XTFY, X0TX0, X0TDX0, X0TFX0,
                                               XTX0, XTDX0, XTFX0, X0TY, X0TDY,
                                               X0TFY, np.log(snr) * 2, l_idx,
                                               n_C, n_T, n_V, n_run, n_X0,
                                               idx_param_fitU, n_C)[0],
        param0_fitU,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec),
                      rtol=1e-5), 'gradient of fitU incorrect'

    # We test the gradient of _fitV wrt to log(SNR^2) assuming no GP prior.
    X0TAX0, XTAX0, X0TAY, X0TAX0_i, \
        XTAcorrX, XTAcorrY, YTAcorrY, LTXTAcorrY, XTAcorrXL, LTXTAcorrXL = \
        brsa._precompute_ar1_quad_forms(XTY, XTDY, XTFY,

assert np.shape(X0TAX0) == (n_V, n_X0, n_X0), (
        'Dimension of X0TAX0 is wrong by _precompute_ar1_quad_forms()')
    assert np.shape(XTAX0) == (n_V, n_C, n_X0), (
        'Dimension of XTAX0 is wrong by _precompute_ar1_quad_forms()')
    assert X0TAY.shape == X0TY.shape, (
        'Shape of X0TAX0 is wrong by _precompute_ar1_quad_forms()')
    assert np.all(np.isfinite(X0TAX0_i)), (
        'Inverse of X0TAX0 includes NaN or Inf')
    ll0, deriv0 = brsa._loglike_AR1_diagV_fitV(
        param0_fitV[idx_param_fitV['log_SNR2']], X0TAX0, XTAX0, X0TAY,
        X0TAX0_i, XTAcorrX, XTAcorrY, YTAcorrY, LTXTAcorrY, XTAcorrXL,
        LTXTAcorrXL, L_full[l_idx], np.tan(rho1 * np.pi / 2), l_idx, n_C, n_T,
        n_V, n_run, n_X0, idx_param_fitV, n_C, False, False)
    vec = np.zeros(np.size(param0_fitV[idx_param_fitV['log_SNR2']]))
    vec[idx_param_fitV['log_SNR2'][0]] = 1
    dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_diagV_fitV(x, X0TAX0, XTAX0, X0TAY,
                                               X0TAX0_i, XTAcorrX, XTAcorrY,
                                               YTAcorrY, LTXTAcorrY, XTAcorrXL,
                                               LTXTAcorrXL, L_full[l_idx],
                                               np.tan(rho1 * np.pi / 2),
                                               l_idx, n_C, n_T, n_V, n_run,
                                               n_X0, idx_param_fitV, n_C,
                                               False, False)[0],
        param0_fitV[idx_param_fitV['log_SNR2']],
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitV wrt log(SNR2) incorrect for model without GP')

    # We test the gradient of _fitV wrt to log(SNR^2) assuming GP prior.
    ll0, deriv0 = brsa._loglike_AR1_diagV_fitV(
        param0_fitV, X0TAX0, XTAX0, X0TAY, X0TAX0_i, XTAcorrX, XTAcorrY,

i]) * s[:, None, None]**2 * a[:, None, None]
        sXTAcorrY[i] = np.dot(design_mat[i].T, Y[i]) * \
            s[:, None, None] * a[:, None, None]

    # test if the gradients are correct
    print(log_fixed_terms)
    ll0, deriv0 = gbrsa._sum_loglike_marginalized(L_vec, s2XTAcorrX,
                                                  YTAcorrY_diag, sXTAcorrY,
                                                  half_log_det_X0TAX0,
                                                  log_weights, log_fixed_terms,
                                                  l_idx, n_C, n_T, n_V, n_X0,
                                                  n_grid, rank=None)
    # We test the gradient to the Cholesky factor
    vec = np.random.randn(np.size(L_vec))
    vec = vec / np.linalg.norm(vec)
    dd = nd.directionaldiff(
        lambda x: gbrsa._sum_loglike_marginalized(x, s2XTAcorrX, YTAcorrY_diag,
                                                  sXTAcorrY,
                                                  half_log_det_X0TAX0,
                                                  log_weights, log_fixed_terms,
                                                  l_idx, n_C, n_T, n_V, n_X0,
                                                  n_grid, rank=None)[0],
        L_vec,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), 'gradient incorrect'

YTDY_diag, YTFY_diag, XTY, XTDY,
                                               XTFY, X0TX0, X0TDX0, X0TFX0,
                                               XTX0, XTDX0, XTFX0, X0TY, X0TDY,
                                               X0TFY, np.log(snr) * 2, l_idx,
                                               n_C, n_T, n_V, n_run, n_X0,
                                               idx_param_fitU, n_C)[0],
        param0_fitU,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitU wrt to AR(1) coefficient incorrect')

    # We test if the numerical and analytical gradient wrt to the first
    # element of Cholesky factor is correct
    vec = np.zeros(np.size(param0_fitU))
    vec[idx_param_fitU['Cholesky'][0]] = 1
    dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_diagV_fitU(x, XTX, XTDX, XTFX, YTY_diag,
                                               YTDY_diag, YTFY_diag, XTY, XTDY,
                                               XTFY, X0TX0, X0TDX0, X0TFX0,
                                               XTX0, XTDX0, XTFX0, X0TY, X0TDY,
                                               X0TFY, np.log(snr) * 2, l_idx,
                                               n_C, n_T, n_V, n_run, n_X0,
                                               idx_param_fitU, n_C)[0],
        param0_fitU,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitU wrt Cholesky factor incorrect')

    # Test on a random direction
    vec = np.random.randn(np.size(param0_fitU))
    vec = vec / np.linalg.norm(vec)
    dd = nd.directionaldiff(

YTAcorrY, LTXTAcorrY, XTAcorrXL,
                                               LTXTAcorrXL, L_full[l_idx],
                                               np.tan(rho1 * np.pi / 2),
                                               l_idx, n_C, n_T, n_V, n_run,
                                               n_X0, idx_param_fitV, n_C, True,
                                               True, dist2, inten_diff2, 100,
                                               100)[0],
        param0_fitV,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitV wrt intensity length scale of GP incorrect')

    # We test the graident on a random direction
    vec = np.random.randn(np.size(param0_fitV))
    vec = vec / np.linalg.norm(vec)
    dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_diagV_fitV(x, X0TAX0, XTAX0, X0TAY,
                                               X0TAX0_i, XTAcorrX, XTAcorrY,
                                               YTAcorrY, LTXTAcorrY, XTAcorrXL,
                                               LTXTAcorrXL, L_full[l_idx],
                                               np.tan(rho1 * np.pi / 2),
                                               l_idx, n_C, n_T, n_V, n_run,
                                               n_X0, idx_param_fitV, n_C, True,
                                               True, dist2, inten_diff2, 100,
                                               100)[0],
        param0_fitV,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitV incorrect')

YTAcorrY, LTXTAcorrY, XTAcorrXL,
                                               LTXTAcorrXL, L_full[l_idx],
                                               np.tan(rho1 * np.pi / 2),
                                               l_idx, n_C, n_T, n_V, n_run,
                                               n_X0, idx_param_fitV, n_C, True,
                                               True, dist2, inten_diff2, 100,
                                               100)[0],
        param0_fitV,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitV srt log(SNR2) incorrect for model with GP')

    # We test the graident wrt spatial length scale parameter of GP prior
    vec = np.zeros(np.size(param0_fitV))
    vec[idx_param_fitV['c_space']] = 1
    dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_diagV_fitV(x, X0TAX0, XTAX0, X0TAY,
                                               X0TAX0_i, XTAcorrX, XTAcorrY,
                                               YTAcorrY, LTXTAcorrY, XTAcorrXL,
                                               LTXTAcorrXL, L_full[l_idx],
                                               np.tan(rho1 * np.pi / 2),
                                               l_idx, n_C, n_T, n_V, n_run,
                                               n_X0, idx_param_fitV, n_C, True,
                                               True, dist2, inten_diff2, 100,
                                               100)[0],
        param0_fitV,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of fitV wrt spatial length scale of GP incorrect')

    # We test the graident wrt intensity length scale parameter of GP prior
    vec = np.zeros(np.size(param0_fitV))

dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_singpara(x, XTX, XTDX, XTFX, YTY_diag,
                                             YTDY_diag, YTFY_diag, XTY, XTDY,
                                             XTFY, X0TX0, X0TDX0, X0TFX0, XTX0,
                                             XTDX0, XTFX0, X0TY, X0TDY, X0TFY,
                                             l_idx, n_C, n_T, n_V, n_run, n_X0,
                                             idx_param_sing)[0],
        param0_sing,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec), rtol=1e-5), (
        'gradient of singpara wrt Cholesky is incorrect')

    # We test the gradient to a1
    vec = np.zeros(np.size(param0_sing))
    vec[idx_param_sing['a1']] = 1
    dd = nd.directionaldiff(
        lambda x: brsa._loglike_AR1_singpara(x, XTX, XTDX, XTFX, YTY_diag,
                                             YTDY_diag, YTFY_diag, XTY, XTDY,
                                             XTFY, X0TX0, X0TDX0, X0TFX0, XTX0,
                                             XTDX0, XTFX0, X0TY, X0TDY, X0TFY,
                                             l_idx, n_C, n_T, n_V, n_run, n_X0,
                                             idx_param_sing)[0],
        param0_sing,
        vec)
    assert np.isclose(dd, np.dot(deriv0, vec),
                      rtol=1e-5), 'gradient of singpara wrt a1 is incorrect'

    # log likelihood and derivative of the fitU function.
    ll0, deriv0 = brsa._loglike_AR1_diagV_fitU(param0_fitU, XTX, XTDX, XTFX,
                                               YTY_diag, YTDY_diag, YTFY_diag,
                                               XTY, XTDY, XTFY, X0TX0, X0TDX0,
                                               X0TFX0, XTX0, XTDX0, XTFX0,