From 2adad098cf5df467ac7e8168039196a889612c6c Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jason.debacker@gmail.com>
Date: Mon, 5 Oct 2026 10:08:32 -0400
Subject: [PATCH 1/4] expand tau_payroll to 2D

---
 CHANGELOG.md                          |  9 ++++
 docs/book/content/intro/parameters.md |  2 +-
 ogcore/aggregates.py                  | 37 +++++++++++----
 ogcore/default_parameters.json        | 10 ++--
 ogcore/household.py                   | 17 +++++--
 ogcore/parameters.py                  | 16 ++++++-
 ogcore/tax.py                         | 39 ++++++++-------
 tests/test_aggregates.py              | 64 ++++++++++++++++++-------
 tests/test_household.py               | 30 +++++++++---
 tests/test_parameters.py              | 68 +++++++++++++++++++++++++++
 tests/test_tax.py                     |  5 +-
 11 files changed, 233 insertions(+), 64 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 34988e62a..d318d1175 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -6,6 +6,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 
+### Added
+- Allows the linear payroll tax rate `tau_payroll` to vary across lifetime income groups as well as over time. The parameter is now a `T x J` array (`number_dims` changed from 1 to 2 in `default_parameters.json`), extrapolated like `labor_income_tax_noncompliance_rate` and `income_tax_filer`, so a single value (`[[0.1]]`), one value per income group (`[[0.1, ..., 0.2]]`), or a time path (`[[0.1], [0.12]]`) can be entered. Reform JSON files or scripts that previously set `tau_payroll` as a flat list (e.g., `[0.124]`) must now use the nested form (e.g., `[[0.124]]`).
+- Adds a `vary tau_payroll by J` case to `test_FOC_labor` in `test_household.py`, a `tau_payroll` extrapolation test to `test_parameters.py`, and a hand-computed check of `get_payroll_tax_revenue` with group-specific rates in `test_aggregates.py`.
+
+### Updated
+- `tax.income_tax_liab` and `household.FOC_labor` now select `tau_payroll` by `j` (or across all `J` when `j` is `None`) in the same way as the tax noncompliance rates and filer status.
+- `aggregates.get_payroll_tax_revenue` now takes household labor supply `n`, effective labor `e`, and population weights in place of aggregate labor `L`, and computes payroll tax revenue as the population-weighted sum of `tau_payroll[t, j] * w * e * n` across households, which is required when the rate differs across income groups. When `tau_payroll` is zero everywhere, the `frac_tax_payroll` fallback is unchanged.
+- `Specifications.compute_default_params` now collapses a previously extrapolated `T x J` parameter (`labor_income_tax_noncompliance_rate`, `capital_income_tax_noncompliance_rate`, `replacement_rate_adjust`, `income_tax_filer`, `wealth_tax_filer`, `tau_payroll`) back to a single column when `J` changes and the values are the same across groups, so that these parameters need not be re-specified whenever `J` is changed via `update_specifications` (the `T x S x J` demographic and ability parameters still must be).
+
 ### Added in PR #1241
 - Adds a `by_J` keyword argument (default `False`) to the demographic plotting functions in `parameter_plots.py` (`plot_imm_rates`, `plot_mort_rates`, `plot_population`, `plot_fert_rates`, `plot_mort_rates_data`, `plot_omega_fixed`, `plot_imm_fixed`, `plot_population_path`) now that the demographic objects are T x S x J. With `by_J=False`, population distributions are summed across income groups and rates are averaged across income groups using the population weights in `omega`, so one line is plotted for the overall population. With `by_J=True`, one line is plotted per income group and labeled with the j value. The array-based functions take an `omega` (or `omega_list`) argument to supply the weights. Each function checks the dimensions of its inputs and continues to accept objects without an income-group dimension (e.g., `T x S` or `S` arrays).
 
diff --git a/docs/book/content/intro/parameters.md b/docs/book/content/intro/parameters.md
index abe9ee5eb..179919b10 100644
--- a/docs/book/content/intro/parameters.md
+++ b/docs/book/content/intro/parameters.md
@@ -630,7 +630,7 @@ _Out-of-Range Action:_ error
 
 ####  `tau_payroll`  
 _Description:_ Linear payroll tax rate. Set value for base year, click '+' to add value for next year.  All future years not specified are set to last value entered.  
-_Notes:_ Set to zero as default since tax functions include income and payroll taxes.  
+_Notes:_ Set to zero as default since tax functions include income and payroll taxes. Can be specified to vary by type J and over year T.  
 _Value Type:_ float  
 _Valid Range:_ min = 0.0 and max = 0.99  
 _Out-of-Range Action:_ error  
diff --git a/ogcore/aggregates.py b/ogcore/aggregates.py
index a70e7b0de..97732267b 100644
--- a/ogcore/aggregates.py
+++ b/ogcore/aggregates.py
@@ -460,7 +460,7 @@ def revenue(
     # income vs payroll split can be reported; it must NOT be added into
     # the total again (that double-counted payroll revenue -- Issue #1199).
     payroll_tax_revenue = get_payroll_tax_revenue(
-        w, L, iit_payroll_tax_revenue, p, method
+        w, n, e, pop_weights, iit_payroll_tax_revenue, p, method
     )
     business_tax_revenue = tax.get_biz_tax(w, Y, L, K, p_m, p, m, method).sum(
         -1
@@ -489,17 +489,21 @@ def revenue(
     )
 
 
-def get_payroll_tax_revenue(w, L, iit_payroll_tax_revenue, p, method):
+def get_payroll_tax_revenue(
+    w, n, e, pop_weights, iit_payroll_tax_revenue, p, method
+):
     r"""
     Calculate aggregate payroll tax revenue.
 
     How payroll tax revenue is computed depends on how the user has
     chosen to represent payroll taxes in the model.  If payroll taxes
     are included directly through the ``tau_payroll`` parameter, then
-    revenue is the payroll tax rate times aggregate labor income:
+    revenue is the sum across households of the (type-specific) payroll
+    tax rate times household labor income:
 
     .. math::
-        PR_{t} = \tau^{p}_{t}w_{t}L_{t}
+        PR_{t} = \sum_{s=E}^{E+S}\sum_{j=0}^{J}\omega_{s,j,t}
+        \tau^{p}_{j,t}w_{t}e_{j,s,t}n_{j,s,t}
 
     Otherwise, payroll taxes are assumed to be embedded in the estimated
     income and payroll tax functions (the default), and payroll tax
@@ -509,7 +513,12 @@ def get_payroll_tax_revenue(w, L, iit_payroll_tax_revenue, p, method):
 
     Args:
         w (array_like): the real wage rate
-        L (array_like): aggregate labor by industry
+        n (Numpy array): household labor supply (S x J for SS, or
+            T x S x J for TPI)
+        e (Numpy array): effective labor units (S x J for SS, or
+            T x S x J for TPI)
+        pop_weights (Numpy array): population weights for each
+            household (S x J for SS, or T x S x J for TPI)
         iit_payroll_tax_revenue (array_like): aggregate income and
             payroll tax revenue
         p (OG-Core Specifications object): model parameters
@@ -522,14 +531,22 @@ def get_payroll_tax_revenue(w, L, iit_payroll_tax_revenue, p, method):
     """
     if np.any(p.tau_payroll != 0):
         # Payroll taxes are modeled explicitly via tau_payroll, so
-        # revenue is the payroll tax rate times aggregate labor income
-        # (summing labor across industries).
-        L_total = L.sum(-1)
+        # revenue is the (type-specific) payroll tax rate times
+        # household labor income, summed over all households.
         if method == "SS":
-            payroll_tax_revenue = p.tau_payroll[-1] * w * L_total
+            tau_payroll = p.tau_payroll[-1, :].reshape(1, p.J)
+            labor_income = w * e * n
+            payroll_tax_revenue = (
+                tau_payroll * labor_income * pop_weights
+            ).sum()
         else:  # TPI
+            tau_payroll = p.tau_payroll[: p.T, :].reshape(p.T, 1, p.J)
+            w = w[: p.T].reshape(p.T, 1, 1)
+            labor_income = w * e[: p.T, :, :] * n[: p.T, :, :]
             payroll_tax_revenue = (
-                p.tau_payroll[: p.T] * w[: p.T] * L_total[: p.T]
+                (tau_payroll * labor_income * pop_weights[: p.T, :, :])
+                .sum(1)
+                .sum(1)
             )
     else:
         # Payroll taxes are embedded in the income and payroll tax
diff --git a/ogcore/default_parameters.json b/ogcore/default_parameters.json
index 6e4878aa2..c1bffd247 100644
--- a/ogcore/default_parameters.json
+++ b/ogcore/default_parameters.json
@@ -1431,17 +1431,17 @@
         "title": "Linear payroll tax rate",
         "description": "Linear payroll tax rate. Set value for base year, click '+' to add value for next year.  All future years not specified are set to last value entered.",
         "short_description": "Payroll tax rate",
-        "param_notation": "$\\tau^{p}_{t}$",
+        "param_notation": "$\\tau^{p}_{j,t}$",
         "section_1": "Fiscal Policy Parameters",
         "section_2": "Taxes",
-        "notes": "Set to zero as default since tax functions include income and payroll taxes.",
+        "notes": "Set to zero as default since tax functions include income and payroll taxes. Can be specified to vary by type J and over year T.",
         "type": "float",
-        "number_dims": 1,
+        "number_dims": 2,
         "value": [
             {
-                "value": [
+                "value": [[
                     0.0
-                ]
+                ]]
             }
         ],
         "validators": {
diff --git a/ogcore/household.py b/ogcore/household.py
index 88f2da8d8..bc479a164 100644
--- a/ogcore/household.py
+++ b/ogcore/household.py
@@ -124,7 +124,9 @@ def marg_ut_labor(n, chi_n, p):
     ) * (
         (1 - ((eps_low / p.ltilde) ** p.upsilon))
         ** ((1 - p.upsilon) / p.upsilon)
-    ) - (2 * b2 * eps_low)
+    ) - (
+        2 * b2 * eps_low
+    )
     MDU_n[nvec_low] = 2 * b2 * nvec[nvec_low] + b1
     d2 = (
         0.5
@@ -147,7 +149,9 @@ def marg_ut_labor(n, chi_n, p):
     ) * (
         (1 - ((eps_high / p.ltilde) ** p.upsilon))
         ** ((1 - p.upsilon) / p.upsilon)
-    ) - (2 * d2 * eps_high)
+    ) - (
+        2 * d2 * eps_high
+    )
     MDU_n[nvec_high] = 2 * d2 * nvec[nvec_high] + d1
     output = MDU_n * np.squeeze(chi_n)
     output = np.squeeze(output)
@@ -689,45 +693,48 @@ def FOC_labor(
 
     """
     if method == "SS":
-        tau_payroll = p.tau_payroll[-1]
         tau_c = p.tau_c[-1, :]
     elif method == "TPI_scalar":  # for 1st donut ring only
-        tau_payroll = p.tau_payroll[0]
         tau_c = p.tau_c[0, :]
     else:
         length = r.shape[0]
-        tau_payroll = p.tau_payroll[t : t + length]
         tau_c = p.tau_c[t : t + length, :]
     if j is not None:
         if method == "SS":
             tax_noncompliance = p.labor_income_tax_noncompliance_rate[-1, j]
             income_tax_filer = p.income_tax_filer[-1, j]
+            tau_payroll = p.tau_payroll[-1, j]
             e = np.squeeze(p.e[-1, :, j])
         elif method == "TPI_scalar":
             tax_noncompliance = p.labor_income_tax_noncompliance_rate[0, j]
             income_tax_filer = p.income_tax_filer[0, j]
+            tau_payroll = p.tau_payroll[0, j]
             e = np.squeeze(p.e[0, -1, j])
         else:
             tax_noncompliance = p.labor_income_tax_noncompliance_rate[
                 t : t + length, j
             ]
             income_tax_filer = p.income_tax_filer[t : t + length, j]
+            tau_payroll = p.tau_payroll[t : t + length, j]
             e_long = _get_e_long(p)
             e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0))
     else:
         if method == "SS":
             tax_noncompliance = p.labor_income_tax_noncompliance_rate[-1, :]
             income_tax_filer = p.income_tax_filer[-1, :]
+            tau_payroll = p.tau_payroll[-1, :]
             e = np.squeeze(p.e[-1, :, :])
         elif method == "TPI_scalar":
             tax_noncompliance = p.labor_income_tax_noncompliance_rate[0, :]
             income_tax_filer = p.income_tax_filer[0, :]
+            tau_payroll = p.tau_payroll[0, :]
             e = np.squeeze(p.e[0, -1, :])
         else:
             tax_noncompliance = p.labor_income_tax_noncompliance_rate[
                 t : t + length, :
             ]
             income_tax_filer = p.income_tax_filer[t : t + length, :]
+            tau_payroll = p.tau_payroll[t : t + length, :]
             e_long = _get_e_long(p)
             e = np.diag(e_long[t : t + p.S, :, j], max(p.S - length, 0))
     if method == "TPI":
diff --git a/ogcore/parameters.py b/ogcore/parameters.py
index a4a40406a..a6c14cd2f 100644
--- a/ogcore/parameters.py
+++ b/ogcore/parameters.py
@@ -188,7 +188,6 @@ def compute_default_params(self):
             "world_int_rate_annual",
             "adjustment_factor_for_cit_receipts",
             "tau_bq",
-            "tau_payroll",
             "h_wealth",
             "m_wealth",
             "p_wealth",
@@ -233,9 +232,22 @@ def compute_default_params(self):
             "replacement_rate_adjust",
             "income_tax_filer",
             "wealth_tax_filer",
+            "tau_payroll",
         ]
         for item in tp_param_list3:
-            param_in = getattr(self, item)
+            param_in = np.asarray(getattr(self, item))
+            # If J has changed since this parameter was last extrapolated
+            # (e.g., via update_specifications), a previously extrapolated
+            # T+S x J_old array will not match the new J.  When the values
+            # are the same across the old J groups, collapse to one column
+            # so that it can be re-extrapolated to the new J.  Otherwise,
+            # leave it to extrapolate_array to raise the error.
+            if (
+                param_in.ndim == 2
+                and param_in.shape[1] != self.J
+                and np.all(param_in == param_in[:, :1])
+            ):
+                param_in = param_in[:, :1]
             param_out = extrapolate_array(
                 param_in, dims=(self.T + self.S, self.J), item=item
             )
diff --git a/ogcore/tax.py b/ogcore/tax.py
index 3cbd6c084..1d4ce4476 100644
--- a/ogcore/tax.py
+++ b/ogcore/tax.py
@@ -356,6 +356,9 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
                 tax_filer = p.income_tax_filer[t : t + length, j].reshape(
                     length, 1
                 )
+                tau_payroll = p.tau_payroll[t : t + length, j].reshape(
+                    length, 1
+                )
             else:
                 labor_income_tax_compliance_rate = (
                     p.labor_income_tax_noncompliance_rate[t : t + length, j]
@@ -364,6 +367,7 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
                     p.capital_income_tax_noncompliance_rate[t : t + length, j]
                 )
                 tax_filer = p.income_tax_filer[t : t + length, j]
+                tau_payroll = p.tau_payroll[t : t + length, j]
         elif method == "TPI_scalar":
             labor_income_tax_compliance_rate = (
                 p.labor_income_tax_noncompliance_rate[0, j]
@@ -372,6 +376,7 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
                 p.capital_income_tax_noncompliance_rate[0, j]
             )
             tax_filer = p.income_tax_filer[0, j]
+            tau_payroll = p.tau_payroll[0, j]
         else:
             labor_income_tax_compliance_rate = (
                 p.labor_income_tax_noncompliance_rate[-1, j]
@@ -380,6 +385,7 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
                 p.capital_income_tax_noncompliance_rate[-1, j]
             )
             tax_filer = p.income_tax_filer[-1, j]
+            tau_payroll = p.tau_payroll[-1, j]
     else:
         if method == "TPI":
             r = utils.to_timepath_shape(r)
@@ -398,6 +404,18 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
             tax_filer = p.income_tax_filer[t : t + length, :].reshape(
                 length, 1, p.J
             )
+            tau_payroll = p.tau_payroll[t : t + length, :].reshape(
+                length, 1, p.J
+            )
+        elif method == "TPI_scalar":
+            labor_income_tax_compliance_rate = (
+                p.labor_income_tax_noncompliance_rate[0, :]
+            )
+            capital_income_tax_compliance_rate = (
+                p.capital_income_tax_noncompliance_rate[0, :]
+            )
+            tax_filer = p.income_tax_filer[0, :]
+            tau_payroll = p.tau_payroll[0, :]
         else:
             labor_income_tax_compliance_rate = (
                 p.labor_income_tax_noncompliance_rate[-1, :]
@@ -406,6 +424,7 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
                 p.capital_income_tax_noncompliance_rate[-1, :]
             )
             tax_filer = p.income_tax_filer[-1, :]
+            tau_payroll = p.tau_payroll[-1, :]
     income = r * b + w * e * n
     labor_income = w * e * n
 
@@ -426,23 +445,9 @@ def income_tax_liab(r, w, b, n, factor, t, j, method, e, etr_params, p):
         * income
     )
 
-    if method == "SS":
-        T_P = p.tau_payroll[-1] * labor_income
-    elif method == "TPI":
-        length = w.shape[0]
-        if len(b.shape) == 1:
-            T_P = p.tau_payroll[t : t + length] * labor_income
-        elif len(b.shape) == 2:
-            T_P = (
-                p.tau_payroll[t : t + length].reshape(length, 1) * labor_income
-            )
-        else:
-            T_P = (
-                p.tau_payroll[t : t + length].reshape(length, 1, 1)
-                * labor_income
-            )
-    elif method == "TPI_scalar":
-        T_P = p.tau_payroll[0] * labor_income
+    # Linear payroll tax, with a rate that may vary over time and by
+    # lifetime income group j (tau_payroll is T+S x J)
+    T_P = tau_payroll * labor_income
 
     income_payroll_tax_liab = T_I + T_P
 
diff --git a/tests/test_aggregates.py b/tests/test_aggregates.py
index 890dc77af..9429fa782 100644
--- a/tests/test_aggregates.py
+++ b/tests/test_aggregates.py
@@ -1309,7 +1309,7 @@ def test_get_C(c, p, method, expected):
     "eta": (np.ones((20, 2)) / (20 * 2)),
     "lambdas": [0.6, 0.4],
     "tau_bq": [0.17],
-    "tau_payroll": [0.5],
+    "tau_payroll": [[0.5]],
     "h_wealth": [0.1],
     "p_wealth": [0.2],
     "m_wealth": [1.0],
@@ -1379,7 +1379,7 @@ def test_get_C(c, p, method, expected):
     "eta": (np.ones((20, 2)) / (20 * 2)),
     "lambdas": [0.6, 0.4],
     "tau_bq": [0.17],
-    "tau_payroll": [0.5],
+    "tau_payroll": [[0.5]],
     "h_wealth": [0.1],
     "p_wealth": [0.2],
     "m_wealth": [1.0],
@@ -1430,7 +1430,7 @@ def test_get_C(c, p, method, expected):
     "eta": (np.ones((20, 2)) / (20 * 2)),
     "lambdas": [0.6, 0.4],
     "tau_bq": [0.17],
-    "tau_payroll": [0.5],
+    "tau_payroll": [[0.5]],
     "h_wealth": [0.1],
     "p_wealth": [0.2],
     "m_wealth": [1.0],
@@ -1778,33 +1778,65 @@ def test_get_payroll_tax_revenue():
     """
     p = Specifications()
     p.T = 3
+    p.S = 2
+    p.J = 2
     iit_payroll_ss = 10.0
     iit_payroll_tpi = np.array([10.0, 11.0, 12.0])
     w_ss = 1.2
-    L_ss = np.array([2.0, 3.0])  # labor by industry
     w_tpi = np.array([1.0, 1.1, 1.2])
-    L_tpi = np.array([[1.0, 2.0], [1.5, 2.5], [2.0, 3.0]])  # (T, M)
+    # household labor supply, effective labor, and population weights
+    # (S x J for SS and T x S x J for TPI)
+    n_ss = np.array([[0.5, 0.4], [0.3, 0.2]])
+    e_ss = np.array([[1.0, 2.0], [1.5, 2.5]])
+    omega_ss = np.array([[0.3, 0.2], [0.3, 0.2]])
+    n_tpi = np.tile(n_ss.reshape(1, 2, 2), (p.T + 1, 1, 1))
+    n_tpi[1, :, :] *= 2.0
+    e_tpi = np.tile(e_ss.reshape(1, 2, 2), (p.T + 1, 1, 1))
+    omega_tpi = np.tile(omega_ss.reshape(1, 2, 2), (p.T + 1, 1, 1))
 
     # Payroll taxes embedded in the tax functions (tau_payroll == 0)
-    p.tau_payroll = np.zeros(p.T)
+    p.tau_payroll = np.zeros((p.T, p.J))
     p.frac_tax_payroll = np.array([0.5, 0.5, 0.5])
-    pr_ss = aggr.get_payroll_tax_revenue(w_ss, L_ss, iit_payroll_ss, p, "SS")
+    pr_ss = aggr.get_payroll_tax_revenue(
+        w_ss, n_ss, e_ss, omega_ss, iit_payroll_ss, p, "SS"
+    )
     assert np.allclose(pr_ss, 0.5 * iit_payroll_ss)
     pr_tpi = aggr.get_payroll_tax_revenue(
-        w_tpi, L_tpi, iit_payroll_tpi, p, "TPI"
+        w_tpi, n_tpi, e_tpi, omega_tpi, iit_payroll_tpi, p, "TPI"
     )
     assert np.allclose(pr_tpi, 0.5 * iit_payroll_tpi)
 
-    # Payroll taxes modeled explicitly via tau_payroll
-    p.tau_payroll = np.array([0.1, 0.2, 0.3])
-    pr_ss = aggr.get_payroll_tax_revenue(w_ss, L_ss, iit_payroll_ss, p, "SS")
-    assert np.allclose(pr_ss, 0.3 * w_ss * L_ss.sum())
-    pr_tpi = aggr.get_payroll_tax_revenue(
-        w_tpi, L_tpi, iit_payroll_tpi, p, "TPI"
+    # Payroll taxes modeled explicitly via tau_payroll, with the rate
+    # varying over time and across lifetime income groups
+    p.tau_payroll = np.array([[0.1, 0.4], [0.2, 0.5], [0.3, 0.6]])
+    pr_ss = aggr.get_payroll_tax_revenue(
+        w_ss, n_ss, e_ss, omega_ss, iit_payroll_ss, p, "SS"
     )
-    assert np.allclose(
-        pr_tpi, np.array([0.1, 0.2, 0.3]) * w_tpi * L_tpi.sum(-1)
+    # hand computation: sum_{s,j} omega_{s,j} tau_j w e_{s,j} n_{s,j}
+    # using the last period rates, tau = [0.3, 0.6]
+    expected_ss = w_ss * (
+        0.3 * 0.3 * 1.0 * 0.5
+        + 0.3 * 0.3 * 1.5 * 0.3
+        + 0.6 * 0.2 * 2.0 * 0.4
+        + 0.6 * 0.2 * 2.5 * 0.2
+    )
+    assert np.allclose(pr_ss, expected_ss)
+    pr_tpi = aggr.get_payroll_tax_revenue(
+        w_tpi, n_tpi, e_tpi, omega_tpi, iit_payroll_tpi, p, "TPI"
     )
+    expected_tpi = np.zeros(p.T)
+    for t in range(p.T):
+        for s in range(p.S):
+            for j in range(p.J):
+                expected_tpi[t] += (
+                    omega_tpi[t, s, j]
+                    * p.tau_payroll[t, j]
+                    * w_tpi[t]
+                    * e_tpi[t, s, j]
+                    * n_tpi[t, s, j]
+                )
+    assert pr_tpi.shape == (p.T,)
+    assert np.allclose(pr_tpi, expected_tpi)
 
 
 test_data = [
diff --git a/tests/test_household.py b/tests/test_household.py
index db33b51d3..d12a36dfc 100644
--- a/tests/test_household.py
+++ b/tests/test_household.py
@@ -605,7 +605,7 @@ def test_get_cons(model_args, expected):
 p1.h_wealth = np.array([0.1])
 p1.m_wealth = np.array([1.0])
 p1.p_wealth = np.array([0.0])
-p1.tau_payroll = np.array([0.15])
+p1.tau_payroll = np.ones((1, p1.J)) * 0.15
 p1.retire = np.array([2]).astype(int)
 
 test_params_ss = p1
@@ -663,7 +663,7 @@ def test_get_cons(model_args, expected):
 # Define variables/params for test of TPI version
 method_tpi = "TPI"
 test_params_tpi = copy.deepcopy(p1)
-test_params_tpi.tau_payroll = np.array([0.15, 0.15, 0.15])
+test_params_tpi.tau_payroll = np.ones((3, p1.J)) * 0.15
 test_params_tpi.tau_bq = np.array([0.0, 0.0, 0.0])
 test_params_tpi.retire = np.array([2, 2, 2]).astype(int)
 test_params_tpi.h_wealth = np.array([0.1, 0.1, 0.1])
@@ -926,7 +926,7 @@ def test_FOC_savings(model_vars, in_params, expected):
 p1.h_wealth = np.array([0.1])
 p1.m_wealth = np.array([1.0])
 p1.p_wealth = np.array([0.0])
-p1.tau_payroll = np.array([0.15])
+p1.tau_payroll = np.ones((1, p1.J)) * 0.15
 p1.retire = np.array([2]).astype(int)
 theta = np.array([0.1])
 j = 0
@@ -970,7 +970,7 @@ def test_FOC_savings(model_vars, in_params, expected):
 test_params_tpi.h_wealth = np.array([0.1, 0.1, 0.1])
 test_params_tpi.m_wealth = np.array([1.0, 1.0, 1.0])
 test_params_tpi.p_wealth = np.array([0.0, 0.0, 0.0])
-test_params_tpi.tau_payroll = np.array([0.15, 0.15, 0.15])
+test_params_tpi.tau_payroll = np.ones((3, p1.J)) * 0.15
 test_params_tpi.tau_bq = np.array([0.0, 0.0, 0.0])
 r_vec = np.array([0.05, 0.03, 0.04])
 w_vec = np.array([1.2, 0.9, 0.8])
@@ -1017,7 +1017,9 @@ def test_FOC_savings(model_vars, in_params, expected):
 test_params_tau_pay.h_wealth = np.array([0.1, 0.1, 0.1])
 test_params_tau_pay.m_wealth = np.array([1.0, 1.0, 1.0])
 test_params_tau_pay.p_wealth = np.array([0.0, 0.0, 0.0])
-test_params_tau_pay.tau_payroll = np.array([0.11, 0.05, 0.33])
+test_params_tau_pay.tau_payroll = np.tile(
+    np.array([0.11, 0.05, 0.33]).reshape(3, 1), (1, p1.J)
+)
 test_params_tau_pay.tau_bq = np.array([0.0, 0.0, 0.0])
 expected_tau_pay = np.array(
     [
@@ -1027,6 +1029,14 @@ def test_FOC_savings(model_vars, in_params, expected):
     ]
 )
 
+# Define variables/params for test with tau_payroll that varies by J.
+# The j=0 column matches test_params_tau_pay and the other columns are
+# set to a very different rate, so the FOC error for j=0 (the group
+# tested) is unchanged only if the j-specific payroll tax rate is used.
+test_params_tau_pay_J = copy.deepcopy(test_params_tau_pay)
+test_params_tau_pay_J.tau_payroll = np.ones((3, 3)) * 0.99
+test_params_tau_pay_J.tau_payroll[:, 0] = np.array([0.11, 0.05, 0.33])
+
 # create parameter objects with non-zero tax noncompliance
 test_params_ss_noncomply = copy.deepcopy(test_params_ss)
 test_params_ss_noncomply.labor_income_tax_noncompliance_rate = (
@@ -1056,6 +1066,7 @@ def test_FOC_savings(model_vars, in_params, expected):
     (test_vars_ss, test_params_ss, expected_ss),
     (test_vars_tpi, test_params_tpi, expected_tpi),
     (test_vars_tpi, test_params_tau_pay, expected_tau_pay),
+    (test_vars_tpi, test_params_tau_pay_J, expected_tau_pay),
     (test_vars_ss, test_params_ss_noncomply, expected_ss_noncomply),
     (test_vars_tpi, test_params_tpi_noncomply, expected_tpi_noncomply),
 ]
@@ -1064,7 +1075,14 @@ def test_FOC_savings(model_vars, in_params, expected):
 @pytest.mark.parametrize(
     "model_vars,params,expected",
     test_data,
-    ids=["SS", "TPI", "vary tau_payroll", "SS, noncomply", "TPI, noncomply"],
+    ids=[
+        "SS",
+        "TPI",
+        "vary tau_payroll",
+        "vary tau_payroll by J",
+        "SS, noncomply",
+        "TPI, noncomply",
+    ],
 )
 def test_FOC_labor(model_vars, params, expected):
     # Test FOC condition for household's choice of labor supply
diff --git a/tests/test_parameters.py b/tests/test_parameters.py
index f35b5e042..810b2e53f 100644
--- a/tests/test_parameters.py
+++ b/tests/test_parameters.py
@@ -53,6 +53,74 @@ def test_alpha_FA_extended_over_time_path():
     assert specs.alpha_FA[-1] == 0.03
 
 
+def test_tau_payroll_extended_over_T_and_J():
+    # tau_payroll may vary over the time path and across lifetime income
+    # groups, so it should be extrapolated to T+S x J from a scalar, a
+    # vector of length J, or a (partial) time path.
+    specs = Specifications()
+    assert specs.tau_payroll.shape == (specs.T + specs.S, specs.J)
+    assert np.all(specs.tau_payroll == 0.0)
+    # single value applies to all groups and periods
+    specs.update_specifications({"tau_payroll": [[0.1]]})
+    assert specs.tau_payroll.shape == (specs.T + specs.S, specs.J)
+    assert np.allclose(specs.tau_payroll, 0.1)
+    # one value per income group, constant over time
+    by_J = [[0.01 * (j + 1) for j in range(specs.J)]]
+    specs.update_specifications({"tau_payroll": by_J})
+    assert specs.tau_payroll.shape == (specs.T + specs.S, specs.J)
+    assert np.allclose(specs.tau_payroll[0, :], by_J[0])
+    assert np.allclose(specs.tau_payroll[-1, :], by_J[0])
+    # time path common across groups, last value carried forward
+    specs.update_specifications({"tau_payroll": [[0.1], [0.2], [0.3]]})
+    assert specs.tau_payroll.shape == (specs.T + specs.S, specs.J)
+    assert np.allclose(specs.tau_payroll[:3, 0], [0.1, 0.2, 0.3])
+    assert np.allclose(specs.tau_payroll[3:, :], 0.3)
+    assert np.allclose(specs.tau_payroll[1, :], 0.2)
+
+
+def test_TxJ_params_collapse_when_J_changes():
+    # A previously extrapolated T+S x J array that is uniform across J
+    # should be re-extrapolated to the new J when J is changed, rather
+    # than raising an error, so the T x J parameters need not be
+    # re-specified when changing J (the T x S x J demographic and
+    # ability parameters still must be).
+    specs = Specifications()
+    specs.update_specifications({"tau_payroll": [[0.1]]})
+    assert specs.tau_payroll.shape[1] == 7
+    rho_vec = np.zeros((1, 20))
+    rho_vec[0, -1] = 1.0
+    new_J_params = {
+        "T": 30,
+        "S": 20,
+        "J": 2,
+        "chi_n": np.ones(2),
+        "e": np.ones((20, 2)),
+        "eta": (np.ones((20, 2)) / (20 * 2)),
+        "lambdas": [0.6, 0.4],
+        "omega": (np.ones((30, 20)) / 20).reshape(30, 20, 1)
+        * np.array([0.6, 0.4]).tolist(),
+        "omega_SS": (np.ones(20) / 20).reshape(20, 1)
+        * np.array([0.6, 0.4]).tolist(),
+        "omega_S_preTP": (np.ones(20) / 20).reshape(20, 1)
+        * np.array([0.6, 0.4]).tolist(),
+        "imm_rates": np.zeros((30, 20, 2)).tolist(),
+        "imm_rates_preTP": np.zeros((20, 2)).tolist(),
+        "rho": np.tile(rho_vec.reshape(1, 20, 1), (30, 1, 2)).tolist(),
+        "rho_preTP": np.tile(rho_vec.reshape(20, 1), (1, 2)).tolist(),
+    }
+    specs.update_specifications(new_J_params)
+    for item in [
+        "tau_payroll",
+        "labor_income_tax_noncompliance_rate",
+        "capital_income_tax_noncompliance_rate",
+        "replacement_rate_adjust",
+        "income_tax_filer",
+        "wealth_tax_filer",
+    ]:
+        assert getattr(specs, item).shape == (specs.T + specs.S, 2), item
+    assert np.allclose(specs.tau_payroll, 0.1)
+
+
 rho_array = np.zeros((4, 3, 7))
 rho_array[:, -1, :] = 1.0
 param_updates1 = {
diff --git a/tests/test_tax.py b/tests/test_tax.py
index 98bc990dd..d54ac8196 100644
--- a/tests/test_tax.py
+++ b/tests/test_tax.py
@@ -968,7 +968,7 @@ def test_get_biz_tax(w, Y, L, K, p_m, p, m, method, expected):
 p.h_wealth = np.ones(p.T + p.S) * 1
 p.p_wealth = np.ones(p.T + p.S) * 2
 p.m_wealth = np.ones(p.T + p.S) * 3
-p.tau_payroll = np.ones(p.T + p.S) * 0.15
+p.tau_payroll = np.ones((p.T + p.S, p.J)) * 0.15
 p.tau_bq = np.ones(p.T + p.S) * 0.1
 p.retire = (np.ones(p.T + p.S) * 2).astype(int)
 p1 = copy.deepcopy(p)
@@ -990,6 +990,7 @@ def test_get_biz_tax(w, Y, L, K, p_m, p, m, method, expected):
 p5.income_tax_filer = np.ones((p5.T, p5.J))
 p5.wealth_tax_filer = np.ones((p5.T, p5.J))
 p5.replacement_rate_adjust = np.ones((p5.T, p5.J))
+p5.tau_payroll = np.ones((p5.T + p5.S, p5.J)) * 0.15
 p5.lambdas = np.array([0.65, 0.35])
 # set variables and other parameters for each case
 r1 = 0.04
@@ -1252,7 +1253,7 @@ def test_get_biz_tax(w, Y, L, K, p_m, p, m, method, expected):
 p_u.h_wealth = np.ones(p_u.T + p_u.S) * 1
 p_u.p_wealth = np.ones(p_u.T + p_u.S) * 2
 p_u.m_wealth = np.ones(p_u.T + p_u.S) * 3
-p_u.tau_payroll = np.ones(p_u.T + p_u.S) * 0.15
+p_u.tau_payroll = np.ones((p_u.T + p_u.S, p_u.J)) * 0.15
 p_u.tau_bq = np.ones(p_u.T + p_u.S) * 0.1
 p_u.retire = (np.ones(p_u.T + p_u.S) * 2).astype(int)
 

From 27b87e81cb0ae0a8fe948908f85aa27ffb6edf3d Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jason.debacker@gmail.com>
Date: Mon, 5 Oct 2026 11:02:18 -0400
Subject: [PATCH 2/4] bump version and update changelog

---
 CHANGELOG.md       | 10 ++++------
 ogcore/__init__.py |  2 +-
 pyproject.toml     |  2 +-
 3 files changed, 6 insertions(+), 8 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d318d1175..1bd739940 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,21 +5,18 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.20.4] - 2026-10-05
 
 ### Added
 - Allows the linear payroll tax rate `tau_payroll` to vary across lifetime income groups as well as over time. The parameter is now a `T x J` array (`number_dims` changed from 1 to 2 in `default_parameters.json`), extrapolated like `labor_income_tax_noncompliance_rate` and `income_tax_filer`, so a single value (`[[0.1]]`), one value per income group (`[[0.1, ..., 0.2]]`), or a time path (`[[0.1], [0.12]]`) can be entered. Reform JSON files or scripts that previously set `tau_payroll` as a flat list (e.g., `[0.124]`) must now use the nested form (e.g., `[[0.124]]`).
 - Adds a `vary tau_payroll by J` case to `test_FOC_labor` in `test_household.py`, a `tau_payroll` extrapolation test to `test_parameters.py`, and a hand-computed check of `get_payroll_tax_revenue` with group-specific rates in `test_aggregates.py`.
+- - Adds a `by_J` keyword argument (default `False`) to the demographic plotting functions in `parameter_plots.py` (`plot_imm_rates`, `plot_mort_rates`, `plot_population`, `plot_fert_rates`, `plot_mort_rates_data`, `plot_omega_fixed`, `plot_imm_fixed`, `plot_population_path`) now that the demographic objects are T x S x J. With `by_J=False`, population distributions are summed across income groups and rates are averaged across income groups using the population weights in `omega`, so one line is plotted for the overall population. With `by_J=True`, one line is plotted per income group and labeled with the j value. The array-based functions take an `omega` (or `omega_list`) argument to supply the weights. Each function checks the dimensions of its inputs and continues to accept objects without an income-group dimension (e.g., `T x S` or `S` arrays).
 
 ### Updated
 - `tax.income_tax_liab` and `household.FOC_labor` now select `tau_payroll` by `j` (or across all `J` when `j` is `None`) in the same way as the tax noncompliance rates and filer status.
 - `aggregates.get_payroll_tax_revenue` now takes household labor supply `n`, effective labor `e`, and population weights in place of aggregate labor `L`, and computes payroll tax revenue as the population-weighted sum of `tau_payroll[t, j] * w * e * n` across households, which is required when the rate differs across income groups. When `tau_payroll` is zero everywhere, the `frac_tax_payroll` fallback is unchanged.
 - `Specifications.compute_default_params` now collapses a previously extrapolated `T x J` parameter (`labor_income_tax_noncompliance_rate`, `capital_income_tax_noncompliance_rate`, `replacement_rate_adjust`, `income_tax_filer`, `wealth_tax_filer`, `tau_payroll`) back to a single column when `J` changes and the values are the same across groups, so that these parameters need not be re-specified whenever `J` is changed via `update_specifications` (the `T x S x J` demographic and ability parameters still must be).
-
-### Added in PR #1241
-- Adds a `by_J` keyword argument (default `False`) to the demographic plotting functions in `parameter_plots.py` (`plot_imm_rates`, `plot_mort_rates`, `plot_population`, `plot_fert_rates`, `plot_mort_rates_data`, `plot_omega_fixed`, `plot_imm_fixed`, `plot_population_path`) now that the demographic objects are T x S x J. With `by_J=False`, population distributions are summed across income groups and rates are averaged across income groups using the population weights in `omega`, so one line is plotted for the overall population. With `by_J=True`, one line is plotted per income group and labeled with the j value. The array-based functions take an `omega` (or `omega_list`) argument to supply the weights. Each function checks the dimensions of its inputs and continues to accept objects without an income-group dimension (e.g., `T x S` or `S` arrays).
-
-### Updated in PR #1241
-- Fixes the population weighting in `plot_mort_rates`, which multiplied mortality rates by the unnormalized population shares rather than taking a weighted average across income groups.
+- - Fixes the population weighting in `plot_mort_rates`, which multiplied mortality rates by the unnormalized population shares rather than taking a weighted average across income groups.
 - Fixes the year indexing in `plot_imm_rates`, `plot_mort_rates_data`, and `plot_population_path`, which indexed years as `start_year - year` rather than `year - start_year`.
 
 ## [0.20.3] - 2026-10-01
@@ -738,6 +735,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Any earlier versions of OG-USA can be found in the [`OG-Core`](https://github.com/PSLmodels/OG-Core) repository [release history](https://github.com/PSLmodels/OG-Core/releases) from [v.0.6.4](https://github.com/PSLmodels/OG-Core/releases/tag/v0.6.4) (Jul. 20, 2021) or earlier.
 
 
+[0.20.4]: https://github.com/PSLmodels/OG-Core/compare/v0.20.3...v0.20.4
 [0.20.3]: https://github.com/PSLmodels/OG-Core/compare/v0.20.2...v0.20.3
 [0.20.2]: https://github.com/PSLmodels/OG-Core/compare/v0.20.1...v0.20.2
 [0.20.1]: https://github.com/PSLmodels/OG-Core/compare/v0.20.0...v0.20.1
diff --git a/ogcore/__init__.py b/ogcore/__init__.py
index 8a38dbbf4..dcd045ef1 100644
--- a/ogcore/__init__.py
+++ b/ogcore/__init__.py
@@ -21,4 +21,4 @@
 from ogcore.txfunc import *  # noqa: F403
 from ogcore.utils import *  # noqa: F403
 
-__version__ = "0.20.3"
+__version__ = "0.20.4"
diff --git a/pyproject.toml b/pyproject.toml
index e9c94d8da..c81d06338 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "ogcore"
-version = "0.20.3"
+version = "0.20.4"
 authors = [
     {name = "Jason DeBacker and Richard W. Evans"},
 ]

From 495d38da96f48c943c4fe2e58d2d15e853be2682 Mon Sep 17 00:00:00 2001
From: Jason DeBacker <jason.debacker@gmail.com>
Date: Mon, 5 Oct 2026 11:08:30 -0400
Subject: [PATCH 3/4] format and update uv lock

---
 ogcore/household.py | 8 ++------
 uv.lock             | 4 ++--
 2 files changed, 4 insertions(+), 8 deletions(-)

diff --git a/ogcore/household.py b/ogcore/household.py
index bc479a164..31ef0b4b0 100644
--- a/ogcore/household.py
+++ b/ogcore/household.py
@@ -124,9 +124,7 @@ def marg_ut_labor(n, chi_n, p):
     ) * (
         (1 - ((eps_low / p.ltilde) ** p.upsilon))
         ** ((1 - p.upsilon) / p.upsilon)
-    ) - (
-        2 * b2 * eps_low
-    )
+    ) - (2 * b2 * eps_low)
     MDU_n[nvec_low] = 2 * b2 * nvec[nvec_low] + b1
     d2 = (
         0.5
@@ -149,9 +147,7 @@ def marg_ut_labor(n, chi_n, p):
     ) * (
         (1 - ((eps_high / p.ltilde) ** p.upsilon))
         ** ((1 - p.upsilon) / p.upsilon)
-    ) - (
-        2 * d2 * eps_high
-    )
+    ) - (2 * d2 * eps_high)
     MDU_n[nvec_high] = 2 * d2 * nvec[nvec_high] + d1
     output = MDU_n * np.squeeze(chi_n)
     output = np.squeeze(output)
diff --git a/uv.lock b/uv.lock
index 25ccd8332..5990da103 100644
--- a/uv.lock
+++ b/uv.lock
@@ -1521,7 +1521,7 @@ wheels = [
 
 [[package]]
 name = "ogcore"
-version = "0.20.3"
+version = "0.20.4"
 source = { editable = "." }
 dependencies = [
     { name = "dask" },
@@ -1707,7 +1707,7 @@ name = "pexpect"
 version = "4.9.0"
 source = { registry = "https://pypi.org/simple" }
 dependencies = [
-    { name = "ptyprocess" },
+    { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" },
 ]
 sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" }
 wheels = [

From 222ed85a493fdd0a874a96587d960f5e4cbe1b94 Mon Sep 17 00:00:00 2001
From: Richard Evans <rickecon@gmail.com>
Date: Mon, 5 Oct 2026 14:07:47 -0600
Subject: [PATCH 4/4] Update CHANGELOG for payroll tax and plotting changes

Updated the CHANGELOG to reflect new features, updates, and fixes related to payroll tax rate variations and demographic plotting functions.
---
 CHANGELOG.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1bd739940..6adeb3f13 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,13 +10,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 - Allows the linear payroll tax rate `tau_payroll` to vary across lifetime income groups as well as over time. The parameter is now a `T x J` array (`number_dims` changed from 1 to 2 in `default_parameters.json`), extrapolated like `labor_income_tax_noncompliance_rate` and `income_tax_filer`, so a single value (`[[0.1]]`), one value per income group (`[[0.1, ..., 0.2]]`), or a time path (`[[0.1], [0.12]]`) can be entered. Reform JSON files or scripts that previously set `tau_payroll` as a flat list (e.g., `[0.124]`) must now use the nested form (e.g., `[[0.124]]`).
 - Adds a `vary tau_payroll by J` case to `test_FOC_labor` in `test_household.py`, a `tau_payroll` extrapolation test to `test_parameters.py`, and a hand-computed check of `get_payroll_tax_revenue` with group-specific rates in `test_aggregates.py`.
-- - Adds a `by_J` keyword argument (default `False`) to the demographic plotting functions in `parameter_plots.py` (`plot_imm_rates`, `plot_mort_rates`, `plot_population`, `plot_fert_rates`, `plot_mort_rates_data`, `plot_omega_fixed`, `plot_imm_fixed`, `plot_population_path`) now that the demographic objects are T x S x J. With `by_J=False`, population distributions are summed across income groups and rates are averaged across income groups using the population weights in `omega`, so one line is plotted for the overall population. With `by_J=True`, one line is plotted per income group and labeled with the j value. The array-based functions take an `omega` (or `omega_list`) argument to supply the weights. Each function checks the dimensions of its inputs and continues to accept objects without an income-group dimension (e.g., `T x S` or `S` arrays).
+- Adds a `by_J` keyword argument (default `False`) to the demographic plotting functions in `parameter_plots.py` (`plot_imm_rates`, `plot_mort_rates`, `plot_population`, `plot_fert_rates`, `plot_mort_rates_data`, `plot_omega_fixed`, `plot_imm_fixed`, `plot_population_path`) now that the demographic objects are T x S x J. With `by_J=False`, population distributions are summed across income groups and rates are averaged across income groups using the population weights in `omega`, so one line is plotted for the overall population. With `by_J=True`, one line is plotted per income group and labeled with the j value. The array-based functions take an `omega` (or `omega_list`) argument to supply the weights. Each function checks the dimensions of its inputs and continues to accept objects without an income-group dimension (e.g., `T x S` or `S` arrays).
 
 ### Updated
 - `tax.income_tax_liab` and `household.FOC_labor` now select `tau_payroll` by `j` (or across all `J` when `j` is `None`) in the same way as the tax noncompliance rates and filer status.
 - `aggregates.get_payroll_tax_revenue` now takes household labor supply `n`, effective labor `e`, and population weights in place of aggregate labor `L`, and computes payroll tax revenue as the population-weighted sum of `tau_payroll[t, j] * w * e * n` across households, which is required when the rate differs across income groups. When `tau_payroll` is zero everywhere, the `frac_tax_payroll` fallback is unchanged.
 - `Specifications.compute_default_params` now collapses a previously extrapolated `T x J` parameter (`labor_income_tax_noncompliance_rate`, `capital_income_tax_noncompliance_rate`, `replacement_rate_adjust`, `income_tax_filer`, `wealth_tax_filer`, `tau_payroll`) back to a single column when `J` changes and the values are the same across groups, so that these parameters need not be re-specified whenever `J` is changed via `update_specifications` (the `T x S x J` demographic and ability parameters still must be).
-- - Fixes the population weighting in `plot_mort_rates`, which multiplied mortality rates by the unnormalized population shares rather than taking a weighted average across income groups.
+- Fixes the population weighting in `plot_mort_rates`, which multiplied mortality rates by the unnormalized population shares rather than taking a weighted average across income groups.
 - Fixes the year indexing in `plot_imm_rates`, `plot_mort_rates_data`, and `plot_population_path`, which indexed years as `start_year - year` rather than `year - start_year`.
 
 ## [0.20.3] - 2026-10-01