Type Ia supernovae (SNe Ia) are bright stellar explosions that are used to measure distances to galaxies and thus to trace the cosmic expansion history, constrain cosmological parameters, and help us understand the nature of the Universe. Supernova distances are estimated from their variations in brightness over time and wavelength (light curves or spectral energy distributions - SEDs). However, interstellar dust in the galaxy of the supernova also affects its brightness and can cause a systematic error in the distance estimates if not properly accounted for. While conventional supernova analyses focus on optical light, a wider wavelength range including near-infrared (NIR) light, should help constrain the effect of dust and improve the precision and accuracy of supernova distances. The objective of this project is to develop BayeSN, the next-generation data-driven probabilistic model for SN Ia spectral energy distributions, jointly leveraging optical and NIR data to determine the optimal distance estimates for current and future supernova cosmology surveys, such as the Vera C. Rubin Observatory’s Legacy Survey of Space and Time (LSST) to ensure principled, accurate, and optimal cosmological analyses.