Skip to content

matfree.stochtrace

matfree.stochtrace

Stochastic estimation of traces, diagonals, and more.

matfree.stochtrace.estimator_leave_one_out(integrand: Callable, /, sampler: Callable) -> Callable

Construct a leave-one-out stochastic estimator.

Parameters:

Name Type Description Default
integrand Callable

An integrand that accepts (matvec, samples, *parameters) where samples has shape (num, n). For example, the return-value of leave_one_out_xtrace.

 required
sampler Callable

The sample function, e.g. the return-value of sampler_normal or sampler_signs.

 required

Returns:

Type Description
estimate

A function estimate(matvec, key, *parameters) -> result. This function can be compiled, vectorised, differentiated, or looped over as the user desires.
Source code in matfree/stochtrace.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def estimator_leave_one_out(integrand: Callable, /, sampler: Callable) -> Callable:
    """Construct a leave-one-out stochastic estimator.

    Parameters
    ----------
    integrand
        An integrand that accepts ``(matvec, samples, *parameters)`` where
        ``samples`` has shape ``(num, n)``. For example, the return-value of
        [leave_one_out_xtrace][matfree.stochtrace.leave_one_out_xtrace].
    sampler
        The sample function, e.g. the return-value of
        [sampler_normal][matfree.stochtrace.sampler_normal] or
        [sampler_signs][matfree.stochtrace.sampler_signs].

    Returns
    -------
    estimate
        A function ``estimate(matvec, key, *parameters) -> result``.
        This function can be compiled, vectorised, differentiated,
        or looped over as the user desires.

    """

    def estimate(matvec, key, *parameters):
        samples = sampler(key)
        Qs = integrand(matvec, samples, *parameters)
        return tree.tree_map(lambda s: np.mean(s, axis=0), Qs)

    return estimate

matfree.stochtrace.estimator_leave_one_out_mean_and_sem(integrand: Callable, /, sampler: Callable) -> Callable

Construct a LOO estimator that returns mean and standard error.

Like estimator_leave_one_out, but returns (mean, sem) where sem = std(loo_estimates) / sqrt(num_samples). The LOO integrand produces one estimate per leave-one-out, so their standard deviation is a natural uncertainty measure.

Parameters:

Name Type Description Default
integrand Callable

Any integrand compatible with estimator_leave_one_out.

 required
sampler Callable

The sample function.

 required

Returns:

Type Description
estimate

A function that returns (mean, sem), both scalars (or arrays with the same shape as a single LOO estimate).
Source code in matfree/stochtrace.py 
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
def estimator_leave_one_out_mean_and_sem(
    integrand: Callable, /, sampler: Callable
) -> Callable:
    """Construct a LOO estimator that returns mean and standard error.

    Like [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out],
    but returns ``(mean, sem)`` where ``sem = std(loo_estimates) / sqrt(num_samples)``.
    The LOO integrand produces one estimate per leave-one-out, so their
    standard deviation is a natural uncertainty measure.

    Parameters
    ----------
    integrand
        Any integrand compatible with
        [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out].
    sampler
        The sample function.

    Returns
    -------
    estimate
        A function that returns ``(mean, sem)``, both scalars (or arrays
        with the same shape as a single LOO estimate).
    """

    def estimate(matvec, key, *parameters):
        samples = sampler(key)
        Qs = integrand(matvec, samples, *parameters)
        n_samples = tree.tree_leaves(Qs)[0].shape[0]
        mean = tree.tree_map(lambda s: np.mean(s, axis=0), Qs)
        sem = tree.tree_map(lambda s: np.std(s, axis=0) / np.sqrt(n_samples), Qs)
        return mean, sem

    return estimate

matfree.stochtrace.estimator_monte_carlo(integrand: Callable, /, sampler: Callable) -> Callable

Construct a stochastic trace-/diagonal-estimator.

Parameters:

Name Type Description Default
integrand Callable

An integrand function with signature integrand(matvec, vec, *parameters), where vec is a single sample vector (contrast with estimator_leave_one_out, which passes the full sample batch). Use any of the monte_carlo_* constructors, e.g. monte_carlo_trace, monte_carlo_diagonal, monte_carlo_frobeniusnorm_squared, monte_carlo_rownorms_squared, or any of the monte_carlo_funm_* functions from matfree.funm.

 required
sampler Callable

The sample function. See below for recommendations.

 required

Returns:

Type Description
estimate

A function that maps a random key to an estimate. This function can be compiled, vectorised, differentiated, or looped over as the user desires.
Notes

The statistical efficiency of the estimator for a given sampler depends on properties of the operator, but we can provide some general advice. For an n-dimensional operator (see references): - n > O(100), use sampler_signs. - n < O(100), use sampler_signs if the operator is known to be diagonal-dominant or sampler_sphere otherwise. - If the operator is complex-valued, pass a complex dtype to the sampler to approximately double the efficiency.

References
Source code in matfree/stochtrace.py 
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
def estimator_monte_carlo(integrand: Callable, /, sampler: Callable) -> Callable:
    """Construct a stochastic trace-/diagonal-estimator.

    Parameters
    ----------
    integrand
        An integrand function with signature ``integrand(matvec, vec, *parameters)``,
        where ``vec`` is a single sample vector (contrast with
        [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out],
        which passes the full sample batch).
        Use any of the ``monte_carlo_*`` constructors, e.g.
        [monte_carlo_trace][matfree.stochtrace.monte_carlo_trace],
        [monte_carlo_diagonal][matfree.stochtrace.monte_carlo_diagonal],
        [monte_carlo_frobeniusnorm_squared][matfree.stochtrace.monte_carlo_frobeniusnorm_squared],
        [monte_carlo_rownorms_squared][matfree.stochtrace.monte_carlo_rownorms_squared],
        or any of the ``monte_carlo_funm_*`` functions from [matfree.funm][matfree.funm].
    sampler
        The sample function. See below for recommendations.

    Returns
    -------
    estimate
        A function that maps a random key to an estimate.
        This function can be compiled, vectorised, differentiated,
        or looped over as the user desires.

    Notes
    -----
    The statistical efficiency of the estimator for a given sampler depends on properties
    of the operator, but we can provide some general advice. For an `n`-dimensional operator (see references):
    - `n > O(100)`, use [sampler_signs][matfree.stochtrace.sampler_signs].
    - `n < O(100)`, use [sampler_signs][matfree.stochtrace.sampler_signs] if the operator is known to be diagonal-dominant or [sampler_sphere][matfree.stochtrace.sampler_sphere] otherwise.
    - If the operator is complex-valued, pass a complex dtype to the sampler to approximately double the efficiency.

    References
    ----------
    - Epperly, E. (2023). [Stochastic trace estimation](https://www.ethanepperly.com/index.php/2023/01/26/stochastic-trace-estimation/).
    - Epperly, E. (2024). [Don't use Gaussians in stochastic trace estimation](https://www.ethanepperly.com/index.php/2024/01/28/dont-use-gaussians-in-stochastic-trace-estimation/).
    """

    def estimate(matvecs, key, *parameters):
        samples = sampler(key)
        Qs = func.vmap(lambda vec: integrand(matvecs, vec, *parameters))(samples)
        return tree.tree_map(lambda s: np.mean(s, axis=0), Qs)

    return estimate

matfree.stochtrace.estimator_monte_carlo_mean_and_sem(integrand: Callable, /, sampler: Callable) -> Callable

Construct a stochastic estimator that returns mean and standard error.

Like estimator_monte_carlo, but returns (mean, sem) where sem = std(samples) / sqrt(num_samples) is the standard error of the mean -- the direct uncertainty on the estimate. The number of samples is already encoded in the sampler, so the caller does not need to track it separately.

Parameters:

Name Type Description Default
integrand Callable

Any integrand compatible with estimator_monte_carlo.

 required
sampler Callable

The sample function. See estimator_monte_carlo for recommendations.

 required

Returns:

Type Description
estimate

A function that returns (mean, sem), both with the same PyTree structure as the integrand output.
Source code in matfree/stochtrace.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
def estimator_monte_carlo_mean_and_sem(
    integrand: Callable, /, sampler: Callable
) -> Callable:
    """Construct a stochastic estimator that returns mean and standard error.

    Like [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo],
    but returns ``(mean, sem)`` where ``sem = std(samples) / sqrt(num_samples)``
    is the standard error of the mean -- the direct uncertainty on the estimate.
    The number of samples is already encoded in the sampler,
    so the caller does not need to track it separately.

    Parameters
    ----------
    integrand
        Any integrand compatible with
        [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo].
    sampler
        The sample function. See [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo]
        for recommendations.

    Returns
    -------
    estimate
        A function that returns ``(mean, sem)``, both with the same
        PyTree structure as the integrand output.
    """

    def estimate(matvecs, key, *parameters):
        samples = sampler(key)
        Qs = func.vmap(lambda vec: integrand(matvecs, vec, *parameters))(samples)
        mean = tree.tree_map(lambda s: np.mean(s, axis=0), Qs)
        sem = tree.tree_map(lambda s: np.std(s, axis=0) / np.sqrt(s.shape[0]), Qs)
        return mean, sem

    return estimate

matfree.stochtrace.leave_one_out_xdiag() -> Callable

Construct an integrand for estimating the diagonal using the XDiag algorithm (Epperly et al. 2024).

Returns:

Type Description
integrand

An integrand function compatible with estimator_leave_one_out whose input has the signature (matvec, samples, *params) and whose output is a pytree with each leaf having shape (num_samples, n_k), giving one diagonal estimate per leave-one-out sample.
Notes

The number of samples must be less than or equal to the dimension of the operator.

The sum of the diagonal estimate over all entries is an unbiased estimate of the trace but generally has higher variance than the estimate produced by leave_one_out_xnystrace.

References
  • Epperly EN, Tropp JA, Webber RJ (2024). XTrace: Making the most of every sample in stochastic trace estimation. SIAM J Matrix Anal A. 45.1: 1-23. doi: 10.1137/23M1548323 arXiv: 2301.07825
  • Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis. arXiv: 2512.15929
Source code in matfree/stochtrace.py 
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
def leave_one_out_xdiag() -> Callable:
    """Construct an integrand for estimating the diagonal using the XDiag algorithm (Epperly et al. 2024).

    Returns
    -------
    integrand
        An integrand function compatible with
        [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out]
        whose input has the signature ``(matvec, samples, *params)`` and whose output is a
        pytree with each leaf having shape ``(num_samples, n_k)``, giving one diagonal
        estimate per leave-one-out sample.

    Notes
    -----
    The number of samples must be less than or equal to the dimension of the operator.

    The sum of the diagonal estimate over all entries is an unbiased estimate of the trace but
    generally has higher variance than the estimate produced by
    [leave_one_out_xnystrace][matfree.stochtrace.leave_one_out_xnystrace].

    References
    ----------
    - Epperly EN, Tropp JA, Webber RJ (2024). XTrace: Making the most of every sample in stochastic trace estimation.
        SIAM J Matrix Anal A. 45.1: 1-23.
        doi: [10.1137/23M1548323](https://doi.org/10.1137/23M1548323)
        arXiv: [2301.07825](https://arxiv.org/abs/2301.07825)
    - Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis.
        arXiv: [2512.15929](https://arxiv.org/abs/2512.15929)
    """

    def integrand(matvec, samples, *params):
        sample0 = tree.tree_map(lambda s: s[0], samples)
        _, unflatten = tree.ravel_pytree(sample0)

        Omega = func.vmap(lambda s: tree.ravel_pytree(s)[0])(samples).T
        n, num_samples = Omega.shape

        if num_samples > n:
            raise ValueError(_error_num_samples(num_samples, maxval=n, minval=1))

        def matvec_flat(v):
            return tree.ravel_pytree(matvec(unflatten(v), *params))[0]

        if 2 * num_samples >= n:
            B_mat = _materialize_operator(matvec_flat, Omega[:, 0])
            diag_B = linalg.diagonal(B_mat)
            return func.vmap(unflatten)(
                np.ones((num_samples, 1), dtype=diag_B.dtype) * diag_B
            )

        matvec_flat_adjoint = func.linear_adjoint(matvec_flat, Omega[:, 0])

        Y = func.vmap(matvec_flat, in_axes=-1, out_axes=-1)(Omega)
        Q, R = linalg.qr_reduced(Y)
        (Z,) = func.vmap(matvec_flat_adjoint, in_axes=-1, out_axes=-1)(Q)

        def _diag_exact():
            diag_B = func.vmap(linalg.vdot, in_axes=0)(Z, Q)
            return np.ones(num_samples, dtype=diag_B.dtype) * diag_B[:, None]

        def _diag_estimate():
            S = _qr_leave_one_out_factor(R)
            QS = Q @ S
            S_vd_R = func.vmap(linalg.vdot, in_axes=1)(S, R)

            diag_B_hat = func.vmap(linalg.vdot, in_axes=0)(Z, Q)
            diag_B_hat_loo = diag_B_hat[:, None] - QS * (Z @ S).conj()
            diag_residual_loo = QS * S_vd_R * Omega.conj()
            return diag_B_hat_loo + diag_residual_loo

        Y_rank = np.sum(np.abs(linalg.diagonal(R)) > np.finfo_eps(R.dtype))

        diag_loo = control_flow.cond(Y_rank < num_samples, _diag_exact, _diag_estimate)
        return func.vmap(unflatten)(diag_loo.T)

    return integrand

matfree.stochtrace.leave_one_out_xnysdiag(*, nystrom: Callable[[Callable, Array], tuple[Array, Array, Array]] | None = None) -> Callable

Construct an integrand for estimating the diagonal of a positive semi-definite operator using the XNysDiag algorithm (Epperly et al. 2025).

Parameters:

Name Type Description Default
nystrom Callable[[Callable, Array], tuple[Array, Array, Array]] | None

A callable with signature (matvec_flat, Omega) -> (nystrom_left, downdate, shift). Usually the return value of nystrom_shifted_cholesky or nystrom_eigh (default: nystrom_eigh).

 None

Returns:

Type Description
integrand

An integrand function compatible with estimator_leave_one_out whose input has the signature (matvec, samples, *params) and whose output is a pytree with each leaf having shape (num_samples, n_k), giving one diagonal estimate per leave-one-out sample. The matvec must be a positive semi-definite operator.
Notes

The number of samples must be less than or equal to the dimension of the operator. The output diagonal is real-valued (PSD operators have real diagonal).

The sum of the diagonal estimate over all entries equals the corresponding leave_one_out_xnystrace trace estimate exactly for the same operator and samples (when apply_resphering=False).

References
  • Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis. arXiv: 2512.15929
Source code in matfree/stochtrace.py 
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
def leave_one_out_xnysdiag(
    *, nystrom: Callable[[Callable, Array], tuple[Array, Array, Array]] | None = None
) -> Callable:
    """Construct an integrand for estimating the diagonal of a positive semi-definite operator using the XNysDiag algorithm (Epperly et al. 2025).

    Parameters
    ----------
    nystrom
        A callable with signature ``(matvec_flat, Omega) -> (nystrom_left, downdate, shift)``.
        Usually the return value of
        [`nystrom_shifted_cholesky`][matfree.stochtrace.nystrom_shifted_cholesky]
        or [`nystrom_eigh`][matfree.stochtrace.nystrom_eigh] (default: `nystrom_eigh`).

    Returns
    -------
    integrand
        An integrand function compatible with
        [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out]
        whose input has the signature ``(matvec, samples, *params)`` and whose output is a
        pytree with each leaf having shape ``(num_samples, n_k)``, giving one diagonal
        estimate per leave-one-out sample.
        The `matvec` must be a positive semi-definite operator.

    Notes
    -----
    The number of samples must be less than or equal to the dimension of the operator.
    The output diagonal is real-valued (PSD operators have real diagonal).

    The sum of the diagonal estimate over all entries equals the corresponding
    [leave_one_out_xnystrace][matfree.stochtrace.leave_one_out_xnystrace]
    trace estimate exactly for the same operator and samples (when ``apply_resphering=False``).

    References
    ----------
    - Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis.
        arXiv: [2512.15929](https://arxiv.org/abs/2512.15929)
    """
    if nystrom is None:
        nystrom = nystrom_eigh()

    def integrand(matvec, samples, *params):
        sample0 = tree.tree_map(lambda s: s[0], samples)
        _, unflatten = tree.ravel_pytree(sample0)

        Omega = func.vmap(lambda s: tree.ravel_pytree(s)[0])(samples).T
        n, num_samples = Omega.shape

        if num_samples > n:
            raise ValueError(_error_num_samples(num_samples, maxval=n, minval=1))

        def matvec_flat(v):
            return tree.ravel_pytree(matvec(unflatten(v), *params))[0]

        if num_samples == n:
            B_mat = _materialize_operator(matvec_flat, Omega[:, 0])
            diag_B = linalg.diagonal(B_mat)
            diag_all = np.ones((num_samples, 1), dtype=diag_B.dtype) * diag_B
            return tree.tree_map(lambda x: x.real, func.vmap(unflatten)(diag_all))

        F, Z, shift = nystrom(matvec_flat, Omega)
        Z_vd_Omega = func.vmap(linalg.vdot, in_axes=1)(Z, Omega)

        diag_B_hat = np.sum(linalg.abs2(F), axis=1) - shift
        diag_B_hat_loo = diag_B_hat[:, None] - linalg.abs2(Z)
        diag_res_loo = Z * Z_vd_Omega * Omega.conj()
        diag_loo = (diag_B_hat_loo + diag_res_loo).T
        return tree.tree_map(lambda x: x.real, func.vmap(unflatten)(diag_loo))

    return integrand

matfree.stochtrace.leave_one_out_xnystrace(*, nystrom: Callable[[Callable, Array], tuple[Array, Array, Array]] | None = None, apply_resphering: bool = True, qr_r: Callable[[Array], Array] | None = None) -> Callable

Construct an integrand for estimating the trace of a positive semi-definite operator using the XNysTrace algorithm (Epperly et al. 2024).

Parameters:

Name Type Description Default
nystrom Callable[[Callable, Array], tuple[Array, Array, Array]] | None

A callable with signature (matvec_flat, Omega) -> (nystrom_left, downdate, shift), where Omega has shape (n, num_samples), nystrom_left and downdate have shape (n, num_samples), and shift is a scalar. nystrom_left @ nystrom_left.T.conj() approximates the operator (shifted by shift * I), and subtracting outer(downdate[:, i], downdate[:, i].conj()) approximates it leaving out the i-th column of Omega. Usually the return value of nystrom_shifted_cholesky or nystrom_eigh (default: nystrom_eigh).

 None
apply_resphering bool

If True (default), project test vectors onto the range of the residual matrix, reducing the variance of the trace estimate. Requires test vectors drawn from a rotationally invariant distribution (e.g. normal or sphere). See Epperly, 2025 for more details.

 True
qr_r Callable[[Array], Array] | None

A callable that computes the R factor of a QR decomposition, used if apply_resphering is True. If not provided, linalg.qr_r is used.

 None

Returns:

Type Description
integrand

An integrand function compatible with estimator_leave_one_out whose input has the signature (matvec, samples, *params) and whose output is a vector of trace estimates with shape (samples.shape[0],). The matvec must be a positive semi-definite operator. That is, vdot(v, matvec(v)) is real and non-negative for all vectors v, and vdot(x, matvec(y)) = vdot(matvec(x), y) for all vectors x and y.
References
  • Epperly EN, Tropp JA, Webber RJ (2024). XTrace: Making the most of every sample in stochastic trace estimation. SIAM J Matrix Anal A. 45.1: 1-23. doi: 10.1137/23M1548323 arXiv: 2301.07825
  • Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis. arXiv: 2512.15929
Source code in matfree/stochtrace.py 
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
def leave_one_out_xnystrace(
    *,
    nystrom: Callable[[Callable, Array], tuple[Array, Array, Array]] | None = None,
    apply_resphering: bool = True,
    qr_r: Callable[[Array], Array] | None = None,
) -> Callable:
    """Construct an integrand for estimating the trace of a positive semi-definite operator using the XNysTrace algorithm (Epperly et al. 2024).

    Parameters
    ----------
    nystrom
        A callable with signature ``(matvec_flat, Omega) -> (nystrom_left, downdate, shift)``,
        where ``Omega`` has shape ``(n, num_samples)``, ``nystrom_left`` and ``downdate``
        have shape ``(n, num_samples)``, and ``shift`` is a scalar.
        ``nystrom_left @ nystrom_left.T.conj()`` approximates the operator (shifted by ``shift * I``),
        and subtracting ``outer(downdate[:, i], downdate[:, i].conj())``
        approximates it leaving out the ``i``-th column of ``Omega``.
        Usually the return value of [`nystrom_shifted_cholesky`][matfree.stochtrace.nystrom_shifted_cholesky]
        or [`nystrom_eigh`][matfree.stochtrace.nystrom_eigh] (default: `nystrom_eigh`).
    apply_resphering
        If ``True`` (default), project test vectors onto the range of the
        residual matrix, reducing the variance of the trace estimate.
        Requires test vectors drawn from a rotationally invariant distribution
        (e.g. normal or sphere). See Epperly, 2025 for more details.
    qr_r
        A callable that computes the R factor of a QR decomposition, used if `apply_resphering` is `True`.
        If not provided, `linalg.qr_r` is used.

    Returns
    -------
    integrand
        An integrand function compatible with `estimator_leave_one_out` whose input
        has the signature ``(matvec, samples, *params)`` and whose output is a vector
        of trace estimates with shape ``(samples.shape[0],)``.
        The `matvec` must be a positive semi-definite operator. That is,
        `vdot(v, matvec(v))` is real and non-negative for all vectors `v`,
        and `vdot(x, matvec(y)) = vdot(matvec(x), y)` for all vectors `x` and `y`.

    References
    ----------
    - Epperly EN, Tropp JA, Webber RJ (2024). XTrace: Making the most of every sample in stochastic trace estimation.
        SIAM J Matrix Anal A. 45.1: 1-23.
        doi: [10.1137/23M1548323](https://doi.org/10.1137/23M1548323)
        arXiv: [2301.07825](https://arxiv.org/abs/2301.07825)
    - Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis.
        arXiv: [2512.15929](https://arxiv.org/abs/2512.15929)
    """
    # NOTE: The paper and thesis use the shifted Nystrom approximation with a
    # Cholesky decomposition, but empirically, this is brittle and fails for
    # many low-rank operators. The eigh-based Nystrom approximation seems to be
    # more robust.
    if nystrom is None:
        nystrom = nystrom_eigh()

    # NOTE: The paper computes R via the QR decomposition, while for efficiency, the
    # thesis uses the upper Cholesky factor of the Gram matrix. We use the QR approach
    # because it may be less brittle and prone to NaNs than Cholesky.
    if qr_r is None:
        qr_r = linalg.qr_r

    def integrand(matvec, samples, *params):
        sample0 = tree.tree_map(lambda s: s[0], samples)
        _, unflatten = tree.ravel_pytree(sample0)

        Omega = func.vmap(lambda s: tree.ravel_pytree(s)[0])(samples).T
        n, num_samples = Omega.shape

        if num_samples > n:
            raise ValueError(_error_num_samples(num_samples, maxval=n, minval=1))

        def matvec_flat(v):
            return tree.ravel_pytree(matvec(unflatten(v), *params))[0]

        if num_samples == n:
            # It's faster and more accurate to compute the trace exactly and deterministically
            # when num_samples == n
            B_mat = _materialize_operator(matvec_flat, Omega[:, 0])
            trace_samples = np.ones((num_samples,), dtype=B_mat.dtype) * linalg.trace(
                B_mat
            )
            return trace_samples.real

        F, Z, shift = nystrom(matvec_flat, Omega)

        if apply_resphering:
            # Ensure T (the R factor) is square
            T = qr_r(Omega)[:num_samples, :num_samples]
            S = _qr_leave_one_out_factor(T)
            # Omega projected onto the subspace spanned by Q_i, i.e. Q from qr(Omega_{-i}) leaving out Omega[:, i]
            X = T - S * func.vmap(linalg.vdot, in_axes=1)(S, T)
            # residual is B - B_hat_{-i}, where B_hat_{-i} approximates B leaving out Omega[:, i]
            rank_residual = n - num_samples + 1
            # squared norm of each sample after projection to the subspace spanned by the residual
            sqnorm_Omega = np.sum(linalg.abs2(Omega), axis=0)
            sqnorm_X = np.sum(linalg.abs2(X), axis=0)
            sqnorm_samples_projected = sqnorm_Omega - sqnorm_X
            sqnorm_samples_projected = np.where(
                sqnorm_samples_projected == 0.0, 1.0, sqnorm_samples_projected
            )
            residual_scale = rank_residual / sqnorm_samples_projected
        else:
            residual_scale = 1.0

        # Compute the trace estimate, correcting for shift in _nystrom_shifted
        tr_B_hat = np.sum(linalg.abs2(F)) - shift * n
        tr_B_hat_loo = tr_B_hat - np.sum(linalg.abs2(Z), axis=0)
        tr_residual_loo = linalg.abs2(func.vmap(linalg.vdot, in_axes=1)(Z, Omega))
        return tr_B_hat_loo + residual_scale * tr_residual_loo

    return integrand

matfree.stochtrace.leave_one_out_xrownorms_squared() -> Callable

Construct an integrand for estimating squared row norms using the XRowNorm algorithm (Epperly, 2025).

Returns:

Type Description
integrand

An integrand function compatible with estimator_leave_one_out whose input has the signature (matvec, samples, *params) and whose output is a pytree with each leaf having shape (num_samples, n_k), giving one squared-row-norm estimate per leave-one-out sample.
Notes

To estimate squared column norms instead, pass the adjoint (i.e. conjugate-transpose-conjugate) of the matvec ( see [func.linear_adjoint][matfree.func.linear_adjoint]).

For normal operators (those that commute with their adjoint), leave_one_out_xsymrownorms_squared, performs only 2/3 of the matvecs per sample compared to this algorithm but is typically less accurate per sample.

References
  • Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis. arXiv: 2512.15929
Source code in matfree/stochtrace.py 
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
def leave_one_out_xrownorms_squared() -> Callable:
    """Construct an integrand for estimating squared row norms using the XRowNorm algorithm (Epperly, 2025).

    Returns
    -------
    integrand
        An integrand function compatible with
        [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out]
        whose input has the signature ``(matvec, samples, *params)`` and whose output is a
        pytree with each leaf having shape ``(num_samples, n_k)``, giving one squared-row-norm
        estimate per leave-one-out sample.

    Notes
    -----
    To estimate squared column norms instead, pass the adjoint (i.e. conjugate-transpose-conjugate) of the matvec (
    see [`func.linear_adjoint`][matfree.func.linear_adjoint]).

    For normal operators (those that commute with their adjoint),
    [leave_one_out_xsymrownorms_squared][matfree.stochtrace.leave_one_out_xsymrownorms_squared],
    performs only 2/3 of the matvecs per sample compared to this algorithm but is typically less
    accurate per sample.

    References
    ----------
    - Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis.
        arXiv: [2512.15929](https://arxiv.org/abs/2512.15929)
    """
    return _leave_one_out_rownorms_squared(iterate_subspace=True)

matfree.stochtrace.leave_one_out_xsymrownorms_squared() -> Callable

Construct an integrand for estimating squared row norms of a normal operator using the XSymRowNorm algorithm (Epperly, 2025).

Returns:

Type Description
integrand

An integrand function compatible with estimator_leave_one_out whose input has the signature (matvec, samples, *params) and whose output is a pytree with each leaf having shape (num_samples, n_k), giving one squared-row-norm estimate per leave-one-out sample. The input and output pytree structures are assumed to be identical. The matvec must be a normal operator, i.e. it must commute with its adjoint (conjugate-transpose-conjugate) up to numerical precision.
Notes

For normal operators, the row and column norms are equal, so this algorithm can also be used to estimate squared column norms.

For general (non-normal) operators, use leave_one_out_xrownorms_squared.

References
  • Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis. arXiv: 2512.15929
Source code in matfree/stochtrace.py 
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
def leave_one_out_xsymrownorms_squared() -> Callable:
    """Construct an integrand for estimating squared row norms of a normal operator using the XSymRowNorm algorithm (Epperly, 2025).

    Returns
    -------
    integrand
        An integrand function compatible with
        [estimator_leave_one_out][matfree.stochtrace.estimator_leave_one_out]
        whose input has the signature ``(matvec, samples, *params)`` and whose output is a
        pytree with each leaf having shape ``(num_samples, n_k)``, giving one squared-row-norm
        estimate per leave-one-out sample.
        The input and output pytree structures are assumed to be identical.
        The ``matvec`` must be a normal operator, i.e. it must commute with its adjoint
        (conjugate-transpose-conjugate) up to numerical precision.

    Notes
    -----
    For normal operators, the row and column norms are equal, so this algorithm can also be used to estimate squared column norms.

    For general (non-normal) operators, use
    [leave_one_out_xrownorms_squared][matfree.stochtrace.leave_one_out_xrownorms_squared].

    References
    ----------
    - Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis.
        arXiv: [2512.15929](https://arxiv.org/abs/2512.15929)
    """
    return _leave_one_out_rownorms_squared(iterate_subspace=False)

matfree.stochtrace.leave_one_out_xtrace(*, apply_resphering: bool = True) -> Callable

Construct an integrand for estimating the trace using the XTrace algorithm (Epperly et al. 2024).

Parameters:

Name Type Description Default
apply_resphering bool

If True (default), project test vectors onto the range of the residual matrix, reducing the variance of the trace estimate. Requires test vectors drawn from a rotationally invariant distribution (e.g. normal or sphere). See Epperly, 2025 for more details.

 True

Returns:

Type Description
integrand

An integrand function compatible with estimator_leave_one_out whose input has the signature (matvec, samples, *params) and whose output is a vector of trace estimates with one estimate per sample.
Notes

The number of samples must be less than or equal to the dimension of the operator. Additionally, the algorithm assumes that the samples are unique. For low-dimensional operators, samples generated from sampler_signs may violate this assumption, and it is recommended to use a different sampler instead.

References
  • Epperly EN, Tropp JA, Webber RJ (2024). XTrace: Making the most of every sample in stochastic trace estimation. SIAM J Matrix Anal A. 45.1: 1-23. doi: 10.1137/23M1548323 arXiv: 2301.07825
  • Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis. arXiv: 2512.15929
Source code in matfree/stochtrace.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
def leave_one_out_xtrace(*, apply_resphering: bool = True) -> Callable:
    """Construct an integrand for estimating the trace using the XTrace algorithm (Epperly et al. 2024).

    Parameters
    ----------
    apply_resphering
        If ``True`` (default), project test vectors onto the range of the
        residual matrix, reducing the variance of the trace estimate.
        Requires test vectors drawn from a rotationally invariant distribution
        (e.g. normal or sphere). See Epperly, 2025 for more details.

    Returns
    -------
    integrand
        An integrand function compatible with `estimator_leave_one_out` whose input
        has the signature ``(matvec, samples, *params)`` and whose output is a vector
        of trace estimates with one estimate per sample.

    Notes
    -----
    The number of samples must be less than or equal to the dimension of the operator.
    Additionally, the algorithm assumes that the samples are unique. For low-dimensional
    operators, samples generated from `sampler_signs` may violate this assumption, and
    it is recommended to use a different sampler instead.

    References
    ----------
    - Epperly EN, Tropp JA, Webber RJ (2024). XTrace: Making the most of every sample in stochastic trace estimation.
        SIAM J Matrix Anal A. 45.1: 1-23.
        doi: [10.1137/23M1548323](https://doi.org/10.1137/23M1548323)
        arXiv: [2301.07825](https://arxiv.org/abs/2301.07825)
    - Epperly EN (2025). Make the most of what you have: Resource-efficient randomized algorithms for matrix computations. PhD Thesis.
        arXiv: [2512.15929](https://arxiv.org/abs/2512.15929)
    """

    def integrand(matvec, samples, *params):
        sample0 = tree.tree_map(lambda s: s[0], samples)
        _, unflatten = tree.ravel_pytree(sample0)

        Omega = func.vmap(lambda s: tree.ravel_pytree(s)[0])(samples).T
        n, num_samples = Omega.shape

        if num_samples > n:
            raise ValueError(_error_num_samples(num_samples, maxval=n, minval=1))

        def matvec_flat(v):
            return tree.ravel_pytree(matvec(unflatten(v), *params))[0]

        if 2 * num_samples >= n:
            # It's faster, more accurate, and allocates less memory to compute the trace exactly
            # and deterministically on the materialized operator when num_samples >= n/2
            B_mat = _materialize_operator(matvec_flat, Omega[:, 0])
            return np.ones((num_samples,), dtype=B_mat.dtype) * linalg.trace(B_mat)

        matmat = func.vmap(matvec_flat, in_axes=-1, out_axes=-1)

        Y = matmat(Omega)
        Q, R = linalg.qr_reduced(Y)
        Z = matmat(Q)

        def _trace_exact():
            tr_B = linalg.vdot(Q, Z)
            return np.ones((num_samples,), dtype=R.dtype) * tr_B

        def _trace_estimate():
            S = _qr_leave_one_out_factor(R)

            Q_H = Q.T.conj()
            # tr(H) == tr(B_hat), where B_hat = Q @ Q.H @ B is a low-rank approximation to the operator B
            H = Q_H @ Z
            W = Q_H @ Omega
            T = Z.T.conj() @ Omega
            W_vd_S = func.vmap(linalg.vdot, in_axes=1)(W, S)
            # Omega projected onto the subspace spanned by Q_i, i.e. Q formed leaving out Omega[:, i]
            X = W - S * W_vd_S.conj()
            T_vd_X = func.vmap(linalg.vdot, in_axes=1)(T, X)
            X_vd_HX = func.vmap(linalg.vdot, in_axes=1)(X, H @ X)
            S_vd_R = func.vmap(linalg.vdot, in_axes=1)(S, R)

            if apply_resphering:
                # residual is B - B_hat_{-i}, where B_hat_{-i} approximates B leaving out Omega[:, i]
                rank_residual = n - num_samples + 1
                # squared norm of each sample after projection to the subspace spanned by the residual
                sqnorm_Omega = np.sum(linalg.abs2(Omega), axis=0)
                sqnorm_X = np.sum(linalg.abs2(X), axis=0)
                sqnorm_samples_projected = sqnorm_Omega - sqnorm_X
                has_zero_norm = sqnorm_samples_projected == 0.0
                sqnorm_samples_projected = np.where(
                    has_zero_norm, 1.0, sqnorm_samples_projected
                )
                residual_scale = rank_residual / sqnorm_samples_projected
            else:
                residual_scale = 1.0

            tr_B_hat = linalg.trace(H)
            # tr(B_hat) leaving out one sample
            tr_B_hat_loo = tr_B_hat - func.vmap(linalg.vdot, in_axes=1)(S, H @ S)
            # Hutchinson estimate of tr(B - B_hat_{-i}) using samples[i, :] as probes.
            tr_residual_loo = -T_vd_X + X_vd_HX + S_vd_R * W_vd_S
            return tr_B_hat_loo + residual_scale * tr_residual_loo

        Y_rank = np.sum(np.abs(linalg.diagonal(R)) > np.finfo_eps(R.dtype))

        # NOTE: If Y_rank < num_samples, then Y is rank-deficient because either:
        # 1. rank(B) < num_samples, and/or
        # 2. the samples are not unique.
        # This check assumes samples are unique, which can be violated for low n and Rademacher samples.
        # If rank(B) < num_samples, then B_hat=B_hat_{-i}=B, so the residual is zero and tr(B_hat)=tr(B).
        return control_flow.cond(Y_rank < num_samples, _trace_exact, _trace_estimate)

    return integrand

matfree.stochtrace.monte_carlo_diagonal()

Construct the integrand for estimating the diagonal.

Use with estimator_monte_carlo. The result will be an Array or PyTree of Arrays with the same tree-structure as matvec(*args_like) where *args_like is an argument of the sampler.

Source code in matfree/stochtrace.py 
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
def monte_carlo_diagonal():
    """Construct the integrand for estimating the diagonal.

    Use with [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo].
    The result will be an Array or PyTree of Arrays with the same tree-structure as
    ``matvec(*args_like)`` where ``*args_like`` is an argument of the sampler.
    """

    def integrand(matvec, v, *parameters):
        Qv = matvec(v, *parameters)
        v_flat, unflatten = tree.ravel_pytree(v)
        Qv_flat, _unflatten = tree.ravel_pytree(Qv)
        return unflatten(v_flat.conj() * Qv_flat)

    return integrand

matfree.stochtrace.monte_carlo_frobeniusnorm_squared()

Construct the integrand for estimating the squared Frobenius norm.

Use with estimator_monte_carlo.

Source code in matfree/stochtrace.py 
903
904
905
906
907
908
909
910
911
912
913
914
def monte_carlo_frobeniusnorm_squared():
    """Construct the integrand for estimating the squared Frobenius norm.

    Use with [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo].
    """

    def integrand(matvec, vec, *parameters):
        x = matvec(vec, *parameters)
        v_flat, _unflatten = tree.ravel_pytree(x)
        return linalg.inner(v_flat.conj(), v_flat)

    return integrand

matfree.stochtrace.monte_carlo_rownorms_squared()

Construct the integrand for estimating the squared row norms.

Use with estimator_monte_carlo.

Notes

To estimate squared column norms instead, pass the adjoint (i.e. conjugate-transpose-conjugate) of the matvec.

Source code in matfree/stochtrace.py 
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
def monte_carlo_rownorms_squared():
    """Construct the integrand for estimating the squared row norms.

    Use with [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo].

    Notes
    -----
    To estimate squared column norms instead, pass the adjoint (i.e. conjugate-transpose-conjugate) of the matvec.
    """

    def integrand(matvec, v, *parameters):
        Qv = matvec(v, *parameters)
        Qv_flat, unflatten = tree.ravel_pytree(Qv)
        return unflatten(linalg.abs2(Qv_flat))

    return integrand

matfree.stochtrace.monte_carlo_trace()

Construct the integrand for estimating the trace.

Use with estimator_monte_carlo.

Source code in matfree/stochtrace.py 
853
854
855
856
857
858
859
860
861
862
863
864
865
def monte_carlo_trace():
    """Construct the integrand for estimating the trace.

    Use with [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo].
    """

    def integrand(matvec, v, *parameters):
        Qv = matvec(v, *parameters)
        v_flat, _unflatten = tree.ravel_pytree(v)
        Qv_flat, _unflatten = tree.ravel_pytree(Qv)
        return linalg.inner(v_flat.conj(), Qv_flat)

    return integrand

matfree.stochtrace.monte_carlo_trace_and_diagonal()

Construct the integrand for estimating the trace and diagonal jointly.

Use with estimator_monte_carlo.

Source code in matfree/stochtrace.py 
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
def monte_carlo_trace_and_diagonal():
    """Construct the integrand for estimating the trace and diagonal jointly.

    Use with [estimator_monte_carlo][matfree.stochtrace.estimator_monte_carlo].
    """

    def integrand(matvec, v, *parameters):
        Qv = matvec(v, *parameters)
        v_flat, unflatten = tree.ravel_pytree(v)
        Qv_flat, _unflatten = tree.ravel_pytree(Qv)
        trace_form = linalg.inner(v_flat.conj(), Qv_flat)
        diagonal_form = unflatten(v_flat.conj() * Qv_flat)
        return {"trace": trace_form, "diagonal": diagonal_form}

    return integrand

matfree.stochtrace.nystrom_eigh(eigenvalues_rtol: float | None = None, leverage_rtol: float | None = None, symmetrize_input: bool = True)

Construct a Nystrom approximation of an operator using a Hermitian eigendecomposition.

Parameters:

Name Type Description Default
eigenvalues_rtol float | None

A relative tolerance used to determine which eigenvalues are close enough to 0.

 None
leverage_rtol float | None

A relative tolerance used in computing the leverage scores to determine which test vectors are essential.

 None
symmetrize_input bool

If True (default), internally symmetrizes before computing the eigendecomposition.

 True

Returns:

Type Description
nystrom

A function that computes the Nystrom approximation of an operator using a Hermitian eigendecomposition. The function has the signature (matvec_flat, Omega) -> (nystrom_left, downdate, shift), where nystrom_left is a left factor of the Nystrom approximation matrix of shape (n, num_samples), such that nystrom_left @ nystrom_left.T.conj() approximates the operator, downdate is a matrix of shape (n, num_samples) whose columns are downdate vectors for the Nystrom approximation, and shift=0 is the shift used (for common API).
Source code in matfree/stochtrace.py 
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
def nystrom_eigh(
    eigenvalues_rtol: float | None = None,
    leverage_rtol: float | None = None,
    symmetrize_input: bool = True,
):
    """Construct a Nystrom approximation of an operator using a Hermitian eigendecomposition.

    Parameters
    ----------
    eigenvalues_rtol
        A relative tolerance used to determine which eigenvalues are close enough to 0.
    leverage_rtol
        A relative tolerance used in computing the leverage scores to determine which
        test vectors are essential.
    symmetrize_input
        If ``True`` (default), internally symmetrizes before computing the eigendecomposition.

    Returns
    -------
    nystrom
        A function that computes the Nystrom approximation of an operator using a Hermitian eigendecomposition.
        The function has the signature `(matvec_flat, Omega) -> (nystrom_left, downdate, shift)`,
        where `nystrom_left` is a left factor of the Nystrom approximation matrix of shape ``(n, num_samples)``,
        such that `nystrom_left @ nystrom_left.T.conj()` approximates the operator,
        `downdate` is a matrix of shape ``(n, num_samples)`` whose columns are downdate vectors for the Nystrom approximation,
        and `shift=0` is the shift used (for common API).
    """

    def nystrom(matvec_flat, Omega):
        num_samples = Omega.shape[1]
        Y = func.vmap(matvec_flat, in_axes=-1, out_axes=-1)(Omega)
        # select rtol using same heuristic as jax.numpy.linalg.lstsq
        if eigenvalues_rtol is None:
            vals_rtol = np.finfo_eps(Y.dtype) * num_samples
        else:
            vals_rtol = eigenvalues_rtol
        H = Omega.T.conj() @ Y

        # Compute left-square-root of pinv(H)
        if symmetrize_input:
            H = _symmetrize(H)
        H_eigh = linalg.eigh(H)
        vals = H_eigh.eigenvalues
        vecs = H_eigh.eigenvectors
        mask = vals >= vals_rtol * np.abs(vals[-1])
        inv_sqrt_vals = np.where(mask, vals ** (-0.5), 0.0)
        vecs = np.where(mask, vecs, 0.0)
        H_pinv_left_sqrt = vecs * inv_sqrt_vals

        # Compute left-square-root of Nystrom approximation
        nystrom_left = Y @ H_pinv_left_sqrt

        # Compute the leverage scores of each column
        leverage = np.sum(np.abs(vecs) ** 2, axis=1)
        if leverage_rtol is None:
            _leverage_rtol = np.sqrt(np.finfo_eps(leverage.dtype))
        else:
            _leverage_rtol = leverage_rtol
        is_essential = leverage + _leverage_rtol > 1.0

        # Compute downdate Z s.t. B_hat_{-i} = B_hat - outer(Z[:, i], Z[:, i].conj()).
        # Since pinv(P H P') = P pinv(H) P', WLOG take i=k with Hk = H without row/col k.
        # Non-essential k, rank(Hk) = rank(H): B_hat_{-k} = B_hat, so Z[:, k] = 0.
        # Essential k, rank(Hk) = rank(H) - 1: by Albert (1969) Thm. 3,
        #   pinv(H) = [pinv(Hk) 0; 0 0] + a v v'  (v = pinv(H)[:, k], a = 1/pinv(H)[k, k]).
        # So B_hat = Y_{-k} pinv(Hk) Y_{-k}' + a (Yv)(Yv)' = B_hat_{-k} + a (Yv)(Yv)',
        # giving Z[:, k] = sqrt(a) Yv = F L[k, :]' / norm(L[k, :])
        # (F, L are left-sqrt factors of B_hat and pinv(H)).
        # Albert (1969). SIAM J. Appl. Math. 17(2), 434-440. doi:10.1137/0117041
        norms = func.vmap(linalg.vector_norm, in_axes=0)(H_pinv_left_sqrt)
        downdate = (nystrom_left @ H_pinv_left_sqrt.T.conj()) / norms
        downdate = np.where(is_essential, downdate, 0.0)

        return nystrom_left, downdate, np.asarray(0.0).astype(vals.dtype)

    return nystrom

matfree.stochtrace.nystrom_shifted_cholesky(shift: float | None = None, rtol: float | None = None, symmetrize_input: bool = True)

Construct a Nystrom approximation of a shifted operator using a Cholesky decomposition.

Parameters:

Name Type Description Default
shift float | None

A small positive shift to add to the operator to ensure the resulting operator is positive definite for Cholesky decomposition. If not provided, the rtol is used to compute the shift.

 None
rtol float | None

A relative tolerance used in computing the shift.

 None
symmetrize_input bool

If True (default), internally symmetrizes before computing the Cholesky factor.

 True

Returns:

Type Description
nystrom

A function that computes the Nystrom approximation of a shifted operator using a Cholesky decomposition. The function has the signature (matvec_flat, Omega) -> (nystrom_left, downdate, shift), where nystrom_left is a left factor of the Nystrom approximation matrix of shape (n, num_samples), such that nystrom_left @ nystrom_left.T.conj() approximates the operator, downdate is a matrix of shape (n, num_samples) whose columns are downdate vectors for the Nystrom approximation, and shift is the shift used.
Source code in matfree/stochtrace.py 
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
def nystrom_shifted_cholesky(
    shift: float | None = None, rtol: float | None = None, symmetrize_input: bool = True
):
    """Construct a Nystrom approximation of a shifted operator using a Cholesky decomposition.

    Parameters
    ----------
    shift
        A small positive shift to add to the operator to ensure the resulting operator
        is positive definite for Cholesky decomposition.
        If not provided, the `rtol` is used to compute the shift.
    rtol
        A relative tolerance used in computing the shift.
    symmetrize_input
        If ``True`` (default), internally symmetrizes before computing the Cholesky factor.

    Returns
    -------
    nystrom
        A function that computes the Nystrom approximation of a shifted operator using a Cholesky decomposition.
        The function has the signature `(matvec_flat, Omega) -> (nystrom_left, downdate, shift)`,
        where `nystrom_left` is a left factor of the Nystrom approximation matrix of shape ``(n, num_samples)``,
        such that `nystrom_left @ nystrom_left.T.conj()` approximates the operator,
        `downdate` is a matrix of shape ``(n, num_samples)`` whose columns are downdate vectors for the Nystrom approximation,
        and `shift` is the shift used.
    """

    def nystrom(matvec_flat, Omega):
        n = Omega.shape[0]
        Y = func.vmap(matvec_flat, in_axes=-1, out_axes=-1)(Omega)
        Y_norm = linalg.vector_norm(Y)
        if shift is None:
            shift_rtol = np.finfo_eps(Y_norm.dtype) if rtol is None else rtol
            mu = shift_rtol * Y_norm / n**0.5
        else:
            mu = shift
        Y_shifted = Y + mu * Omega
        H = Omega.T.conj() @ Y_shifted

        # Compute left-square-root of inv(H)
        if symmetrize_input:
            H = _symmetrize(H)
        H_cholu = linalg.cholesky(H).T.conj()
        Id = np.eye(H_cholu.shape[0], dtype=H_cholu.dtype)
        H_inv_left_sqrt = linalg.solve_triangular(H_cholu, Id)

        # Compute left-square-root of Nystrom approximation
        nystrom_right = linalg.solve_triangular(H_cholu, Y_shifted.T.conj(), trans=2)
        nystrom_left = nystrom_right.T.conj()

        norms = func.vmap(linalg.vector_norm, in_axes=0)(H_inv_left_sqrt)
        downdate = linalg.solve_triangular(H_cholu, nystrom_right).T.conj()
        downdate = downdate / norms[None, :]

        return nystrom_left, downdate, mu

    return nystrom

matfree.stochtrace.sampler_normal(*args_like, num)

Construct a function that samples from a standard-normal distribution.

Source code in matfree/stochtrace.py 
927
928
929
def sampler_normal(*args_like, num):
    """Construct a function that samples from a standard-normal distribution."""
    return _sampler_from_jax_random(prng.normal, *args_like, num=num)

matfree.stochtrace.sampler_signs(*args_like, num)

Construct a function that samples signs uniformly.

For real dtypes, this samples from a Rademacher distribution (uniformly over {-1, 1}). For complex dtypes, this samples from a Steinhaus distribution on the complex unit circle.

Source code in matfree/stochtrace.py 
932
933
934
935
936
937
def sampler_signs(*args_like, num):
    """Construct a function that samples signs uniformly.

    For real dtypes, this samples from a Rademacher distribution (uniformly over `{-1, 1}`). For complex dtypes, this samples from a Steinhaus distribution on the complex unit circle.
    """
    return _sampler_from_jax_random(_uniform_signs, *args_like, num=num)

matfree.stochtrace.sampler_sphere(*args_like, num)

Construct a function that samples from a unit sphere scaled to have identity covariance.

Source code in matfree/stochtrace.py 
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
def sampler_sphere(*args_like, num):
    """Construct a function that samples from a unit sphere scaled to have identity covariance."""
    x_flat, unflatten = tree.ravel_pytree(*args_like)
    dtype = x_flat.dtype
    rdtype = dtype.type(0).real.dtype
    n = x_flat.shape[0]
    sqrtn = np.sqrt(n).astype(rdtype)

    def sample(key):
        samples = prng.normal(key, shape=(num, n), dtype=dtype)
        return func.vmap(lambda x: unflatten(x * (sqrtn / linalg.vector_norm(x))))(
            samples
        )

    return sample