Pith. sign in

REVIEW 3 major objections 5 minor 17 references

jinns: a JAX Library for Physics-Informed Neural Networks

T0 review · 3 major / 5 minor · reviewed 2026-08-11 · deepseek-v4-flash

Pith's one-line read The paper claims that jinns is the first PINN library built entirely on the JAX ecosystem, and that this native design makes it the fastest library for inverse problems in a five-problem benchmark against three established Python PINN…

desk verdict Solid JAX-native PINN library with real engineering value, but the 'fastest for inverse problems' claim overreaches the paper's own benchmark data. read the letter →

arxiv 2412.14132 v1 pith:LHY2IXLS submitted 2024-12-18 stat.ML cs.LG

classification stat.MLcs.LG
keywords physics-informedneuralnetworksPINNsJAXpythoninverseproblemsmeta-modelingopen-sourcesoftware
verification ladder T0 review T1 audit T2 compute T3 formal

The pith

A machine-rendered reading of the paper's core claim, the machinery that carries it, and where it could break.

The reading

The paper introduces jinns, an open-source Python library for physics-informed neural networks built entirely on the JAX ecosystem. Its central claim is twofold: jinns is the first PINN library to be native to that ecosystem, and in a benchmark against three established Python PINN libraries on five problems, it is the fastest for inverse problems, with speed-ups up to about an order of magnitude. The paper argues this matters because no dedicated JAX-native PINN library existed, and one major comparison library's JAX backend cannot handle inverse problems. jinns is designed around an explicit split between network weights and equation parameters, with a derivative mask controlling which loss terms are differentiated with respect to which parameters, so forward problems, inverse problems, and meta-modeling share one interface.

What carries the argument

The argument rides on jinns' modular architecture rather than on a single mathematical identity. The core objects are the Params class, which separates equation parameters from neural-network weights, and the optional DerivativeKeys mask, which controls how each loss component is differentiated with respect to each parameter set; these two mechanisms let one training loop serve forward, inverse, and meta-modeling problems. The library also uses equinox.Module for transparent PyTree registration, so jax.jit, jax.grad, and jax.vmap compose with user-defined models, and it relies on optax for optimization. For evaluation, the benchmark reproduces problems from the PINNacle suite and from the DeepXDE documentation, comparing training times and relative errors against three other Python PINN libraries on an Nvidia T600 GPU.

What would settle it

Re-run the five benchmark problems on at least two different GPUs with many repeated trials per library, and check the benchmark scripts for identical optimizers, learning-rate schedules, collocation-point sets, and stopping criteria across all libraries; if another library is faster on average for the inverse problems, or if the scripts reveal unequal hyperparameters, the paper's central speed claim is refuted.

Watch

Extended reading notes

Core claim

The paper's central claim is a software design and an associated empirical result: a modular PINN library written in pure JAX, relying on the equinox and optax libraries for models and optimization, can cover forward problems, inverse problems, and meta-model learning in a uniform interface while remaining fast. The design separates the neural network parameters from the differential operator parameters, and lets users specify a derivative mask so the physics loss, boundary terms, or data terms can be differentiated with respect to either set independently. In the reported benchmark, jinns completes inverse problems in the shortest training time among the compared libraries, with gains up to roughly tenfold over the next fastest on problems such as diffusion-reaction and Navier-Stokes parameter identification, at error levels comparable to the other implementations.

Load-bearing premise

The central performance claim rests on the benchmark setup being fair: all hyperparameters truly equal across libraries, and a single run on one Nvidia T600 GPU being representative; if those do not hold, the conclusion that jinns is the fastest library for inverse problems collapses.

Editorial extensions

If this is right

  • Users of the JAX stack can set up PINN forward and inverse problems without leaving the JAX ecosystem, since models, optimizers, and transformations all come from the same set of libraries.
  • Inverse problems that the JAX backend of the main comparison library cannot express become implementable, because jinns is JAX-native rather than a backend shim.
  • Out-of-the-box baselines include SeparablePINN and HyperPINN architectures alongside vanilla MLPs, so architecture comparisons do not require reimplementing them.
  • If the benchmark timings hold, jinns offers up to a tenfold reduction in training time for inverse problems at comparable estimation error, making parameter-identification studies substantially cheaper.

Reading between the lines

Editorial extensions of the paper, not claims the author makes directly.

  • The derivative-mask design points beyond the reported use cases: the same mechanism should make it straightforward to train meta-models that differentiate the loss with respect to equation parameters on a per-term basis, a pattern the paper leaves implicit.
  • Because jinns is pure JAX, it should compose with JAX's parallel and distributed transforms such as pmap and sharding; the paper does not benchmark these, so whether the speed advantage extends to multi-GPU or batch-parallel settings is untested.
  • The benchmark's single-run, single-GPU measurements are not enough to establish a stable ordering across hardware and random seeds; a repeated-run study would tell whether jinns is reliably the fastest rather than fastest in this particular run.
Share X Bluesky LinkedIn Reddit HN

Editorial analysis

A structured set of objections, weighed in public.

Desk editor's note, referee report, and a circularity audit.

Referee Report

3 major / 5 minor

Summary. The paper presents jinns, an open-source Python library for physics-informed neural networks built natively on the JAX ecosystem, with support for forward problems, inverse problems, and meta-modeling. The authors describe the library's modular design (data, parameters, loss, utility modules), its reliance on equinox and optax, and its use of forward and reverse automatic differentiation, including support for SPINN and HyperPINN architectures. They report a benchmark against DeepXDE, PINA, and NVIDIA Modulus on five problems (Burgers1D, NS2d-C, Poisson inverse, Diffusion-Reaction inverse, Navier-Stokes inverse), claiming that jinns is the fastest library for inverse problems and that the code is reproducible via a dedicated Gitlab repository. The paper also includes code listings illustrating the definition of PDE losses in jinns.

Significance. If its claims are supported, jinns fills a genuine gap in the PINN software ecosystem: a JAX-native library that handles inverse problems and meta-modeling, where the paper states that DeepXDE's JAX backend cannot be used to its full extent. The library appears well-structured and thoughtfully designed: it leverages existing JAX standards (equinox, optax), separates equation parameters from neural-network parameters, offers a DerivativeKeys mask for granular differentiation of loss terms, and provides documentation, tutorials, unit tests, and reproducible benchmark code. The code listings and the description of the workflow indicate that the library is usable and extensible. The central performance claim, however, is not currently supported by the paper's own reported data, and the benchmark methodology is not sufficient to establish a global speed advantage. The software contribution itself is valuable, but the paper's headline claim needs to be corrected or substantially re-evidenced.

major comments (3)
  1. [Section 2.5, Tables 1-3] The statement that 'jinns compares favorably to the other libraries, in particular, it is the fastest library for inverse problems' is contradicted by the paper's own Table 1. On the Poisson inverse problem (PInv), jinns takes 149 s while NVIDIA Modulus JAX takes 135 s, and Tables 2 and 3 show jinns with L1/L2 errors of 0.070/0.088 versus Modulus JAX's 0.034/0.043. The 'fastest' claim therefore holds only for some inverse problems (e.g., Diffusion-Reaction-Inv and Navier-Stokes-Inv), not for the inverse-problem category as a whole. The claim should either be restricted to per-problem statements or supported by more comprehensive evidence.
  2. [Appendix C.2] The benchmark timing results are based on a single run on a single Nvidia T600 GPU, with no error bars, no multiple seeds, and no measure of run-to-run variability. Since PINN training involves stochastic initialization and, for Modulus, stochastic collocation-point sampling, a single-run comparison is not a reliable basis for a 'fastest library' claim. The paper should either report repeated runs with means and variances, or soften the conclusion to describe the observed single-run timings.
  3. [Appendix C.2] The claim that 'all hyper-parameters are equal across all libraries' is immediately qualified by the exception that Modulus uses uniformly sampled collocation points rather than regularly spaced points. This exception means that the benchmark does not actually control the sampling scheme across libraries. Since the paper does not enumerate the chosen hyper-parameters (learning rate, network width/depth, optimizer settings, number of iterations, etc.) beyond referring to scripts in the repository, the reader cannot independently verify that the comparison is fair. A table of hyper-parameters per problem and library, or a more explicit discussion of how the Modulus sampling exception could affect timings and accuracy, is needed.
minor comments (5)
  1. [Section 1] The phrase 'backward problems' is used to mean 'inverse problems'; this terminology is non-standard and may confuse readers. Consider using 'inverse' consistently throughout the introduction.
  2. [Section 2.2] There is a typo in 'inhomegeneous'; it should be 'inhomogeneous'.
  3. [Section 2.1, utils module] The word 'litterature' should be 'literature'.
  4. [Appendix B, Code 2] The code listing defines 'source term(xy)' with one argument but calls 'source term(x, params)' with two arguments; this inconsistency should be corrected.
  5. [Appendix C.3, Tables 2-3] The text states that the last two rows are identical in Tables 2 and 3 because they report errors on two scalar parameters; clarifying how L1 and L2 relative errors are computed for scalar parameter estimates would make the table easier to interpret.

Circularity Check

0 steps flagged · score 0.0 of 10

No circularity: this is a software and benchmark report with no derived prediction that reduces to its inputs; the fastest-library overclaim is a correctness concern, not a circularity concern.

full rationale

The paper makes no substantive mathematical derivation whose conclusion is built into its assumptions. Its central claims are (1) jinns is an open-source JAX-native PINN library, (2) it supports forward, inverse, and meta-modeling problems through modular classes, and (3) it compares favorably to DeepXDE, PINA, and Modulus, 'in particular, it is the fastest library for inverse problems' (Section 2.5). None of these claims is a prediction fitted from data or an equation derived from a definition that already contains the result. The benchmark does involve the authors comparing their own library to others, which creates a conflict-of-interest risk, but that is not circular reasoning under the patterns enumerated here: there is no fitted parameter renamed as a prediction, no load-bearing self-citation, and no imported uniqueness theorem. The performance claim is in fact contradicted by the paper's own Table 1 on the PInv problem, where NVIDIA Modulus JAX (135s) beats jinns (149s), and the Appendix C.2 statement that 'all hyper-parameters are equal across all libraries' is immediately qualified by the exception for Modulus sampling. These are internal-validity and reporting weaknesses, not circularity. I therefore find no significant circularity and assign score 0.

Assumptions & free parameters 0 free parameters · 4 assumptions · 0 invented entities

The paper rests on standard PINN methodology and external software, not on new theoretical constructs. No free parameters or invented entities are introduced. The main assumptions are about the validity of the PINN approach, the trustworthiness of benchmark references, and the correctness of the underlying JAX ecosystem.

assumptions (4)
  • domain assumption Minimizing a weighted sum of PDE residuals and boundary/initial condition losses yields an accurate approximation of the true PDE solution.
    The library and benchmarks are built on this PINN principle (Appendix A); if false for the chosen problems, the reported errors lose meaning.
  • domain assumption The reference solutions and problem formulations taken from PINNacle and the DeepXDE documentation are correct and suitable as ground truth.
    Errors in Tables 2 and 3 are computed against these references (Appendix C.1).
  • domain assumption JAX, equinox, and optax behave as documented, and jinns integrates them without subtle bugs affecting the benchmark outcomes.
    The performance and correctness of jinns depend on these external tools (Sections 2.1, 2.2).
  • domain assumption The claim of being 'the first library for PINNs built solely on the JAX ecosystem' is supported by the authors' literature survey.
    Section 1 states 'To the best of our knowledge'; no systematic survey is provided.

how reviews work

0 comments
Cite this review

Pith. "Pith review of jinns: a JAX Library for Physics-Informed Neural Networks." pith.science (2026). https://pith.science/paper/LHY2IXLS

@misc{pith2026241214132,
  author       = {Pith},
  title        = {Pith review of: jinns: a JAX Library for Physics-Informed Neural Networks},
  year         = {2026},
  howpublished = {\url{https://pith.science/paper/LHY2IXLS}},
  note         = {Machine review of arXiv:2412.14132}
}
read the original abstract

jinns is an open-source Python library for physics-informed neural networks, built to tackle both forward and inverse problems, as well as meta-model learning. Rooted in the JAX ecosystem, it provides a versatile framework for efficiently prototyping real-problems, while easily allowing extensions to specific needs. Furthermore, the implementation leverages existing popular JAX libraries such as equinox and optax for model definition and optimisation, bringing a sense of familiarity to the user. Many models are available as baselines, and the documentation provides reference implementations of different use-cases along with step-by-step tutorials for extensions to specific needs. The code is available on Gitlab https://gitlab.com/mia_jinns/jinns.

Figures

Figures reproduced from arXiv: 2412.14132 by the authors.

Figure 1
Figure 1. Typical user workflow for jinns users. 2 Design and implementation 2.1 Organisation of the library jinns is a Python library with compatibility for Python 3.11+ built on-top of the JAX ecosystem. It is focused on user-flexibility and reproducibility. As such, it provides a modular interface allowing to specify the learning problem, defining a loss, a set of training points, and the parameters of interest to be optim… view at source ↗

Discussion (0). Continue with ORCID to comment.

Reference graph

Works this paper leans on

17 extracted references · 11 canonical work pages

  1. [1]

    Araz, Juan Carlos Criado, and Michael Spannwosky

    Jack Y. Araz, Juan Carlos Criado, and Michael Spannwosky. Elvet -- a neural network-based differential equation and variational problem solver, 2021

  2. [2]

    JAX : composable transformations of P ython+ N um P y programs, 2024

    James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake Vander P las, Skye Wanderman- M ilne, and Qiao Zhang. JAX : composable transformations of P ython+ N um P y programs, 2024. URL http://github.com/google/jax

  3. [3]

    Separable physics-informed neural networks

    Junwoo Cho, Seungtae Nam, Hyunmo Yang, Seok-Bae Yun, Youngjoon Hong, and Eunbyung Park. Separable physics-informed neural networks. Advances in Neural Information Processing Systems, 36, 2024

  4. [4]

    Physics-informed neural networks for advanced modeling

    Dario Coscia, Anna Ivagnes, Nicola Demo, and Gianluigi Rozza. Physics-informed neural networks for advanced modeling. Journal of Open Source Software, 8 0 (87): 0 5352, 2023. doi:10.21105/joss.05352. URL https://doi.org/10.21105/joss.05352

  5. [5]

    Scientific machine learning through physics--informed neural networks: Where we are and what’s next

    Salvatore Cuomo, Vincenzo Schiano Di Cola, Fabio Giampaolo, Gianluigi Rozza, Maziar Raissi, and Francesco Piccialli. Scientific machine learning through physics--informed neural networks: Where we are and what’s next. Journal of Scientific Computing, 92 0 (3): 0 88, 2022

  6. [6]

    Systems biology: Identifiability analysis and parameter identification via systems-biology-informed neural networks

    Mitchell Daneker, Zhen Zhang, George Em Karniadakis, and Lu Lu. Systems biology: Identifiability analysis and parameter identification via systems-biology-informed neural networks. In Computational Modeling of Signaling Networks, pages 87--105. Springer, 2023

  7. [7]

    Hyperpinn: Learning parameterized differential equations with physics-informed hypernetworks

    Filipe de Avila Belbute-Peres, Yi-fan Chen, and Fei Sha. Hyperpinn: Learning parameterized differential equations with physics-informed hypernetworks. The symbiosis of deep learning and differential equations, 690, 2021

  8. [8]

    Pinnacle: A comprehensive benchmark of physics-informed neural networks for solving pdes

    Zhongkai Hao, Jiachen Yao, Chang Su, Hang Su, Ziao Wang, Fanzhi Lu, Zeyu Xia, Yichi Zhang, Songming Liu, Lu Lu, et al. Pinnacle: A comprehensive benchmark of physics-informed neural networks for solving pdes. arXiv preprint arXiv:2306.08827, 2023

Show all 17 references
  1. [9]

    Physics-informed machine learning

    George Em Karniadakis, Ioannis G Kevrekidis, Lu Lu, Paris Perdikaris, Sifan Wang, and Liu Yang. Physics-informed machine learning. Nature Reviews Physics, 3 0 (6): 0 422--440, 2021

  2. [10]

    Fourier neural operator for parametric partial differential equations

    Zongyi Li, Nikola Kovachki, Kamyar Azizzadenesheli, Burigede Liu, Kaushik Bhattacharya, Andrew Stuart, and Anima Anandkumar. Fourier neural operator for parametric partial differential equations. arXiv preprint arXiv:2010.08895, 2020

  3. [11]

    DeepXDE : A deep learning library for solving differential equations

    Lu Lu, Xuhui Meng, Zhiping Mao, and George Em Karniadakis. DeepXDE : A deep learning library for solving differential equations. SIAM Review, 63 0 (1): 0 208--228, 2021. doi:10.1137/19M1274067

  4. [12]

    Nvidia modulus, 2023

    NVIDIA. Nvidia modulus, 2023. URL https://github.com/NVIDIA/modulus

  5. [13]

    Idrlnet: A physics-informed neural network library

    Wei Peng, Jun Zhang, Weien Zhou, Xiaoyu Zhao, Wen Yao, and Xiaoqian Chen. Idrlnet: A physics-informed neural network library. 2021

  6. [14]

    Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations

    Maziar Raissi, Paris Perdikaris, and George E Karniadakis. Physics-informed neural networks: A deep learning framework for solving forward and inverse problems involving nonlinear partial differential equations. Journal of Computational physics, 378: 0 686--707, 2019

  7. [15]

    Spatio-temporal ecological models via physics-informed neural networks for studying chronic wasting disease

    Juan Francisco Mandujano Reyes, Ting Fung Ma, Ian P McGahan, Daniel J Storm, Daniel P Walsh, and Jun Zhu. Spatio-temporal ecological models via physics-informed neural networks for studying chronic wasting disease. Spatial Statistics, 62: 0 100850, 2024

  8. [16]

    Pdebench: An extensive benchmark for scientific machine learning

    Makoto Takamoto, Timothy Praditia, Raphael Leiteritz, Daniel MacKinlay, Francesco Alesiani, Dirk Pfl\" u ger, and Mathias Niepert. Pdebench: An extensive benchmark for scientific machine learning. In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh, editors, A...

  9. [17]

    Neuralpde: Automating physics-informed neural networks (pinns) with error approximations, 2021

    Kirill Zubov, Zoe McCarthy, Yingbo Ma, Francesco Calisto, Valerio Pagliarino, Simone Azeglio, Luca Bottero, Emmanuel Luján, Valentin Sulzer, Ashutosh Bharambe, Nand Vinchhi, Kaushik Balakrishnan, Devesh Upadhyay, and Chris Rackauckas. Neuralpde: Automating physics-informed neu...

Pith tools

Reviewed August 11, 2026 · model on record in the stance chip above.