Neural network emulator to constrain the high-$z$ IGM thermal state from Lyman-$\alpha$ forest flux auto-correlation function

We present a neural network emulator to constrain the thermal parameters of the intergalactic medium (IGM) at $\displaystyle{5.4}\le{z}\le{6.0}$ using the Lyman-$\displaystyle\alpha$ (Ly$\displaystyle\alpha$) forest flux auto-correlation function. Our auto-differentiable JAX-based framework accelerates the surrogate model generation process using approximately 100 sparsely sampled Nyx hydrodynamical simulations with varying combinations of thermal parameters, i.e., the temperature at mean density $\displaystyle{T}_{{0}}$, the slope of the temperature$\displaystyle-$density relation $\displaystyle\gamma$, and the mean transmission flux $\displaystyle{\langle}{F}{\rangle}$. We show that this emulator has a typical accuracy of 1.0% across the specified redshift range. Bayesian inference of the IGM thermal parameters, incorporating emulator uncertainty propagation, is further expedited using NumPyro Hamiltonian Monte Carlo. We compare both the inference results and computational cost of our framework with the traditional nearest-neighbor interpolation approach applied to the same set of mock Ly$\alpha$ flux. By examining the credibility contours of the marginalized posteriors for $\displaystyle{T}_{{0}},\gamma,\text{and}{\langle}{F}{\rangle}$ obtained using the emulator, the statistical reliability of measurements is established through inference on 100 realistic mock data sets of the auto-correlation function.