Tensorization over Factorization: Rethinking Permutation Invariant Polynomials in JAX

Permutation invariant polynomials have been the gold standard descriptor for constructing high-precision potential energy surfaces (PESs) for over two decades. The monomial symmetrization approach (MSA) has long served as the dominant engine for this task, primarily by minimizing floating-point operations (FLOPs) through algebraic factorization. While the common sub-expression factorization leads to non-contiguous memory access and explosion of intermediate monomials, causing deep and scattered computational graph for automatic differentiation (AD) engines. To resolve these issues, we introduce JaxPIP, a library implemented in JAX and Equinox, that advocates the paradigm shift from factorization to tensorization. JaxPIP reformulates PIPs evaluation into dense, hardware-friendly tensor operations via log-exp fusion and segmented summation. Benchmarks show that JaxPIP achieves orders-of-magnitude faster JIT compilation and superior throughput compared to existing implementations such as MOLPIPx on both CPUs and GPUs. The tensorized architecture naturally supports massive batching, enabling efficient high-throughput simulations such as path-integral molecular dynamics (PIMD), ring-polymer molecular dynamics (RPMD) and vibrational diffusion Monte Carlo (VibDMC). Our practice confirms that on modern architectures, the cost of redundant FLOPs is tiny compared to the benefits of memory continuity and computational graph simplicity.