arXiv Science⌕ Search
Home/Archive/arXiv

arXiv · 2112.09017

Large Scale Distributed Linear Algebra With Tensor Processing Units

Adam G. M. Lewis·Jackson Beall·Martin Ganahl·Markus Hauru·Shrestha Basu Mallick·Guifre Vidal
View original source ↗Read source document ↗Export metadata ↓

Abstract

We have repurposed Google Tensor Processing Units (TPUs), application-specific chips developed for machine learning, into large-scale dense linear algebra supercomputers. The TPUs' fast inter-core interconnects (ICI)s, physically two-dimensional network topology, and high-bandwidth memory (HBM) permit distributed matrix multiplication algorithms to rapidly become computationally bound. In this regime, the matrix-multiply units (MXU)s dominate the runtime, yielding impressive scaling, performance, and raw size: operating in float32 precision, a full 2048-core pod of third generation TPUs can multiply two matrices with linear size $N= 220= 1 048 576$ in about 2 minutes. Via curated algorithms emphasizing large, single-core matrix multiplications, other tasks in dense linear algebra can similarly scale. As examples, we present (i) QR decomposition; (ii) resolution of linear systems; and (iii) the computation of matrix functions by polynomial iteration, demonstrated by the matrix polar factorization.

Explore related subjects

physics.comp-phcs.DCquant-ph

Keep this discovery

BibTeXRIS

Adam G. M. Lewis, Jackson Beall, Martin Ganahl, Markus Hauru, Shrestha Basu Mallick, Guifre Vidal. 2021-12-16. Large Scale Distributed Linear Algebra With Tensor Processing Units. https://doi.org/10.1073/pnas.2122762119

Cite the original work for its findings. Save a collection to share your selection of sources.

KEEP EXPLORING

Related papers

arXiv2026

A Geometrically Parameterized Quasi-Stationary 3D Model for High-Frequency Induction Tube Welding

A three-dimensional multiphysics finite element framework for the simulation of high-frequency induction welding of tubes is presented. The model couples a time-harmonic magnetic scalar potential formulation with a stabilized quasi-stationary advection-diffusion heat transport equation, enabling accurate prediction of electromagnetic and thermal fields under industrial operating conditions. The framework incorporates parameterized geometry generation and semi-automated, physics-tailored mesh construction and is implemented using the open-source tools GetDP and Gmsh. Validation against measurements from a commercial induction welding line for AISI 304 stainless steel tubes demonstrates good agreement with operating data. The validated model is subsequently applied to investigate the influence of impeder material by comparing a conventional FeNiZnV ferrite with the soft magnetic composite Ferrotron 559H for the induction welding of AISI 304 stainless steel tubes.

physics.comp-ph↗
arXiv2026

Stress-Testing Dynamical and Generative Downscaling Using Subseasonal Extreme Precipitation Forecasts

Coarse spatial resolution limits the ability of subseasonal prediction models to resolve extreme precipitation. Downscaling with either dynamical or deep generative models can overcome this issue, but the comparative performance of these models for extremes across different atmospheric regimes remains poorly understood. In this work, we evaluate the Weather Research and Forecasting (WRF) model against a diffusion-based generative model by downscaling two physically distinct, extreme precipitation events up to lead times of 3 weeks. For a fair comparison with WRF, which can downscale boundary conditions from different driving models without model-specific training, the diffusion model is trained in an unpaired fashion. Both approaches improve upon the raw European Centre for Medium-Range Weather Forecasts forecasts, in comparison to fused rain gauge-radar observations in Switzerland (CombiPrecip), but exhibit regime-dependent strengths. WRF achieves the highest probabilistic skill for a multicell, non-stationary event. Conversely, the diffusion model is more consistent across different performance metrics for the two events, outperforming WRF in a more stationary supercell event. These results demonstrate that explicit dynamical modeling can add value for specific precipitation events for subseasonal lead times, and that generative downscaling adds value more broadly in different situations.

physics.comp-ph↗
arXiv2026

Nonlinear flame describing function and mean shift kinematics of slit flames under combined axial-transverse forcing

This study investigates the nonlinear kinematics of a premixed slit flame using a two-dimensional $G$-equation level-set framework. Results show that combined forcing induces nonlinear saturation in the FDF, characterized by early gain flattening and premature phase drops, which intensify with the transverse forcing amplitude. Kinematic analysis reveals that this geometric nonlinearity manifests as a reduction in the time-averaged flame height, defined as the mean shift. In the quasi-steady limit, this mean shift is analytically quantified via a multivariate asymptotic expansion, where fourth-order terms successfully capture the saturation mechanism at elevated amplitudes. By introducing a scaling parameter to account for transverse dominance, the frequency-dependent decay of the mean shift in the compact limit collapses onto a single master curve, enabling the derivation of a unified theoretical model that integrates this asymptotic response with a second-order low-pass filter. Furthermore, because the mean shift reduces the physical extent of the flame, it alters the wrinkle propagation time. Correcting the Strouhal number using the measured mean shift collapses the dispersed nonlinear FDF curves onto the linear theory prediction. The analysis is further extended to disturbances convected at a finite speed, for which the linear transfer function is derived analytically and the correction with the measured mean shift continues to collapse the nonlinear FDF. These findings establish that the nonlinear FDF behavior under multidimensional forcing is fundamentally governed by the kinematic mean shift, providing a theoretical baseline for decoupling geometric nonlinearities from other thermo-diffusive or hydrodynamic instabilities in turbulent flames.

physics.comp-ph↗