From 59b0a78a3292ef11c9517372bcbc06fa64499032 Mon Sep 17 00:00:00 2001 From: ahfoysal Date: Tue, 7 Jul 2026 12:13:44 +0600 Subject: [PATCH 1/2] docs: document NUTS regularize_mass_matrix --- numpyro/infer/hmc.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py index 3cd69c659..f54683b4d 100644 --- a/numpyro/infer/hmc.py +++ b/numpyro/infer/hmc.py @@ -870,6 +870,9 @@ class NUTS(HMC): during warm-up phase using Dual Averaging scheme. :param bool adapt_mass_matrix: A flag to decide if we want to adapt mass matrix during warm-up phase using Welford scheme. + :param bool regularize_mass_matrix: whether or not to regularize the estimated mass + matrix for numerical stability during warmup phase. Defaults to True. This flag + does not take effect if ``adapt_mass_matrix == False``. :param dense_mass: This flag controls whether mass matrix is dense (i.e. full-rank) or diagonal (defaults to ``dense_mass=False``). To specify a structured mass matrix, users can provide a list of tuples of site names. Each tuple represents From 1582a1f967142efa86737c8b9c14a806f0f9f3fd Mon Sep 17 00:00:00 2001 From: Abu Hossain Foysal Date: Tue, 7 Jul 2026 20:22:18 +0600 Subject: [PATCH 2/2] docs: move regularize_mass_matrix docs to match signature order --- numpyro/infer/hmc.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/numpyro/infer/hmc.py b/numpyro/infer/hmc.py index f54683b4d..d15b290e8 100644 --- a/numpyro/infer/hmc.py +++ b/numpyro/infer/hmc.py @@ -870,9 +870,6 @@ class NUTS(HMC): during warm-up phase using Dual Averaging scheme. :param bool adapt_mass_matrix: A flag to decide if we want to adapt mass matrix during warm-up phase using Welford scheme. - :param bool regularize_mass_matrix: whether or not to regularize the estimated mass - matrix for numerical stability during warmup phase. Defaults to True. This flag - does not take effect if ``adapt_mass_matrix == False``. :param dense_mass: This flag controls whether mass matrix is dense (i.e. full-rank) or diagonal (defaults to ``dense_mass=False``). To specify a structured mass matrix, users can provide a list of tuples of site names. Each tuple represents @@ -911,6 +908,9 @@ class NUTS(HMC): only supports forward-mode differentiation. See `JAX's The Autodiff Cookbook `_ for more information. + :param bool regularize_mass_matrix: whether or not to regularize the estimated mass + matrix for numerical stability during warmup phase. Defaults to True. This flag + does not take effect if ``adapt_mass_matrix == False``. """ def __init__(