# A Typed Tensor Language for Shared-State Federated Computation

arXiv: https://arxiv.org/abs/2605.21103 · HTML: https://mlsanity.com/paper/2605.21103/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 55 references that are in the map.
- First posted on arXiv 20 May 2026.

## Builds on (most cited first, 25 of 25)

- [Communication-Efficient Learning of Deep Networks from Decentralized Data](https://mlsanity.com/paper/1602.05629/) (2016)
- [Advances and Open Problems in Federated Learning](https://mlsanity.com/paper/1912.04977/) (2019)
- [Federated Optimization in Heterogeneous Networks](https://mlsanity.com/paper/1812.06127/) (2018)
- [Towards Federated Learning at Scale: System Design](https://mlsanity.com/paper/1902.01046/) (2019)
- [The Future of Digital Health with Federated Learning](https://mlsanity.com/paper/2003.08119/) (2020)
- [FedBN: Federated Learning on Non-IID Features via Local Batch Normalization](https://mlsanity.com/paper/2102.07623/) (2021)
- [Federated Learning Based on Dynamic Regularization](https://mlsanity.com/paper/2111.04263/) (2021)
- [Flower: A Friendly Federated Learning Research Framework](https://mlsanity.com/paper/2007.14390/) (2020)
- [Federated Learning for Internet of Things: A Comprehensive Survey](https://mlsanity.com/paper/2104.07914/) (2021)
- [Tensor Programs I: Wide Feedforward or Recurrent Neural Networks of Any Architecture are Gaussian Processes](https://mlsanity.com/paper/1910.12478/) (2019)
- [Federated Learning for Open Banking](https://mlsanity.com/paper/2108.10749/) (2021)
- [Federated Learning for Wireless Communications: Motivation, Opportunities and Challenges](https://mlsanity.com/paper/1908.06847/) (2019)
- [NVIDIA FLARE: Federated Learning from Simulation to Real-World](https://mlsanity.com/paper/2210.13291/) (2022)
- [OpenFL: An open-source framework for Federated Learning](https://mlsanity.com/paper/2105.06413/) (2021)
- [Communication Lower Bounds for Statistical Estimation Problems via a Distributed Data Processing Inequality](https://mlsanity.com/paper/1506.07216/) (2015)
- [SystemDS: A Declarative Machine Learning System for the End-to-End Data Science Lifecycle](https://mlsanity.com/paper/1909.02976/) (2019)
- [Applications of Federated Learning in Smart Cities: Recent Advances, Taxonomy, and Open Challenges](https://mlsanity.com/paper/2102.01375/) (2021)
- [Tensor Relational Algebra for Machine Learning System Design](https://mlsanity.com/paper/2009.00524/) (2020)
- [Getting to the Point. Index Sets and Parallelism-Preserving Autodiff for Pointful Array Programming](https://mlsanity.com/paper/2104.05372/) (2021)
- [DrJAX: Scalable and Differentiable MapReduce Primitives in JAX](https://mlsanity.com/paper/2403.07128/) (2024)
- [Empowering Federated Learning for Massive Models with NVIDIA FLARE](https://mlsanity.com/paper/2402.07792/) (2024)
- [FBFL: A Field-Based Coordination Approach for Data Heterogeneity in Federated Learning](https://mlsanity.com/paper/2502.08577/) (2025)
- [Federated Automatic Differentiation](https://mlsanity.com/paper/2301.07806/) (2023)
- [GRACEFUL: A Learned Cost Estimator For UDFs](https://mlsanity.com/paper/2503.23863/) (2025)
- [Domain-Specific Tensor Languages](https://mlsanity.com/paper/2312.02664/) (2023)

## Abstract (arXiv)

Shared-state federated computations combine client-local tensor computation, mergeable aggregation into shared state, and shared-only post-processing. We introduce a typed tensor language for this class of computations. Its two tensor sorts separate client-partitioned data from globally available values, and typing tracks the partitioned axis. A virtual global tensor serves as a semantic reference for centralized evaluation. We show that typed one-round programs factor through shared tensors whose shapes depend on the program but are independent of client and sample counts. The converse applies to typed-realizable factorizations: each encoder component is represented by an allowed aggregation or contraction with its valid merge, and the decoder is shared-only. The construction extends round by round to programs whose persistent state is shared. For a loss supplied with a client-local per-sample gradient expression, summation represents the empirical gradient. This gives typed programs for server-side first-order updates and, with shared linear algebra, curvature-block updates. The language covers federated analytics and FedSGD. General multi-local-step FedAvg and persistent private client state are outside its scope.