<i>traceax</i> : a JAX-based framework for stochastic trace estimation.

ABSTRACT: In many applications, from statistical inference to machine learning, calculating the trace of a matrix is a fundamental operation, yet may be infeasible due to memory constraints. Stochastic trace estimation offers a practical solution by using randomized matrix-vector products to obtain accurate, unbiased estimates without constructing the full matrix in memory. Here, we present traceax , a Python framework for scalable trace estimation that leverages efficient linear operator representations of matrices while supporting automatic differentiation and hardware acceleration. traceax supports state-of-the-art trace estimators and through simulations we recapitulate results demonstrating their high accuracy while significantly reducing runtime and memory usage when compared wit

SUBMITTER: Nahid AA 

PROVIDER: S-EPMC12338724 | biostudies-literature | 2025 Jul

REPOSITORIES: biostudies-literature