This paper introduces JAX-BTE, a GPU-accelerated, differentiable solver for the phonon Boltzmann Transport Equation (BTE) based on differentiable programming. JAX-BTE enables accurate, efficient and differentiable multiscale thermal modeling by leveraging high-performance GPU computing and automatic differentiation. The solver efficiently addresses the high-dimensional and complex integro-differential nature of the phonon BTE, facilitating both forward simulations and data-augmented inverse simulations through end-to-end optimization. Validation is performed across a range of 1D to 3D simulations, including complex FinFET structures, in both forward and inverse settings, demonstrating excellent performance and reliability. JAX-BTE significantly outperforms state-of-the-art BTE solvers in forward simulations and uniquely enables inverse simulations, making it a powerful tool for multiscale thermal analysis and design for semiconductor devices.
Cite
@article{arxiv.2503.23657,
title = {JAX-BTE: A GPU-Accelerated Differentiable Solver for Phonon Boltzmann Transport Equations},
author = {Wenjie Shang and Jiahang Zhou and J. P. Panda and Zhihao Xu and Yi Liu and Pan Du and Jian-Xun Wang and Tengfei Luo},
journal= {arXiv preprint arXiv:2503.23657},
year = {2025}
}