{"doi":"10.1101/2025.07.14.662216","title":"<i>traceax</i> : a JAX-based framework for stochastic trace estimation","abstract":"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 with direct trace computation. As a proof of concept, we implemented a stochastic heritability estimator using traceax requiring only several lines of code. Overall, traceax provides a versatile tool for stochastic trace estimation that can be easily integrated into existing inferential pipelines. traceax is freely available at: https://github.com/mancusolab/traceax","journal":"bioRxiv (Cold Spring Harbor Laboratory)","year":2025,"id":572150,"datarank":0.0,"base_score":0.0,"endowment":0.0,"self_citation_contribution":0.0,"citation_network_contribution":0.0,"self_endowment_contribution":0.0,"citer_contribution":0.0,"corpus_percentile":null,"corpus_rank":null,"citation_count":0,"citer_count":0,"citers_with_citation_signal":0,"citers_with_endowment":0,"datacite_reuse_total":0,"is_dataset":false,"is_dataset_confidence":0.9471,"is_data_producer":false,"deposit_databanks":null,"is_oa":true,"file_count":0,"downloads":0,"has_version_chain":false,"published_date":"2025-01-01","fair_score":null,"fair_percentile":null,"algorithm_id":"datarank_citation_only_1hop_v6","ranking_scope":"data_only","authors":[{"id":1478804,"name":"Linda Serafin","orcid":null,"position":1,"is_corresponding":false},{"id":702,"name":"Nicholas Mancuso","orcid":"0000-0002-9352-5927","position":2,"is_corresponding":false},{"id":1165931,"name":"Abdullah Al Nahid","orcid":"0000-0002-4390-0768","position":0,"is_corresponding":true}],"reference_count":13,"raw_metadata":null,"created_at":"2026-07-19T02:57:23.653298Z","pmid":"40791344","pmcid":null,"fwci":null,"citation_percentile":null,"influential_citations":0,"oa_status":null,"license":null,"views":0,"total_file_size_bytes":0,"version_count":0,"fair_f":null,"fair_a":null,"fair_i":null,"fair_r":null,"fair_zscore":null,"fair_rationale":null,"fair_model":null,"fair_agent_version":null,"fair_fulltext_source":null,"fair_has_llm":null,"fair_computed_at":null,"clinical_trials":[],"software_tools":[],"db_accessions":[],"linked_datasets":[],"topics":[]}