tfp.experimental.fastgp.fast_log_det.log_det_jvp

Jacobian-Vector product for log_det.

tfp.experimental.fastgp.fast_log_det.log_det_jvp(
    primals, tangents, f, num_iters
)

This can be used to provide a jax.custom_jvp for a function with the signature (matrix, preconditioner, probe_vectors, key, *optional_args, num_iters) -> Float.

[num_iters needs to be last in the signature because it needs to be a non-diff argument of the custom.jvp, which means it isn't passed in the primals and we need to put it back in here.]

primals The arguments f was called with.
tangents The differentials of the primals.
f The log det function log_det_jvp is the custom.jvp of.
num_iters The number of iterations to run conjugate gradients for.

The pair (f(primals), df).