def METHOD_NAME(input_dim, hidden_dims, param_dims, skip_connections):
    rng_key, rng_key_perm = random.split(random.PRNGKey(0))
    perm = random.permutation(rng_key_perm, np.arange(input_dim))
    arn_init, arn = AutoregressiveNN(
        input_dim,
        hidden_dims,
        param_dims=param_dims,
        skip_connections=skip_connections,
        permutation=perm,
    )
    batch_size = 4
    input_shape = (batch_size, input_dim)
    _, init_params = arn_init(rng_key, input_shape)
    output = arn(init_params, np.random.rand(*input_shape))
    if param_dims == [1]:
        assert output.shape == (batch_size, input_dim)
        jac = jacfwd(lambda x: arn(init_params, x))(np.random.rand(input_dim))
    elif param_dims == [1, 1]:
        assert output[0].shape == (batch_size, input_dim)
        assert output[1].shape == (batch_size, input_dim)
        jac = jacfwd(lambda x: arn(init_params, x)[0])(np.random.rand(input_dim))
    elif param_dims == [2]:
        assert output.shape == (2, batch_size, input_dim)
        jac = jacfwd(lambda x: arn(init_params, x))(np.random.rand(input_dim))
    elif param_dims == [2, 3]:
        assert output[0].shape == (2, batch_size, input_dim)
        assert output[1].shape == (3, batch_size, input_dim)
        jac = jacfwd(lambda x: arn(init_params, x)[0])(np.random.rand(input_dim))
    # permute jacobian as necessary
    permuted_jac = np.zeros(jac.shape)
    for j in range(input_dim):
        for k in range(input_dim):
            permuted_jac[..., j, k] = jac[..., perm[j], perm[k]]
    # make sure jacobians are triangular
    assert np.sum(np.abs(np.triu(permuted_jac))) == 0.0