I need to use `pyro.ops.gaussian.sequential_gaussian_filter_sample` (and similar Gaussian utilities) with near-singular transition matrices over long sequences. Currently, the internal Cholesky decomposition occasionally fails because the precision matrix is numerically non-positive definite, even when the model is theoretically valid. This forces me to add large manual jitter that distorts the results.

**Acceptance criteria:**

- All Cholesky decompositions performed internally by `pyro.ops.gaussian`, `pyro.distributions.hmm`, and related modules must succeed without raising `torch.linalg.cholesky` errors, even when the input matrix is extremely ill‑conditioned (e.g., eigenvalues barely above machine epsilon).
- The library should automatically add a small amount of jitter to the matrix diagonal, proportional to the matrix’s magnitude and the floating‑point epsilon, so that precision is maintained for well‑conditioned matrices and only minimal perturbation is applied when needed.
- The jitter must be applied **in every call site** that previously used a raw Cholesky: in `pyro.ops.gaussian.Gaussian` (sampling, conditioning, marginalizing, event_logsumexp), in `gaussian_tensordot`, and in the HMM classes (`GaussianHMM`, `GammaGaussianHMM`). The user‑facing `safe_cholesky` function in `pyro.ops.tensor_utils` should be the single point of implementation, replacing the existing `cholesky` function.
- For scalar (1×1) precision matrices, the transform should clamp to `torch.finfo(dtype).tiny` to avoid zero or negative values.
- The fix must not break gradient flow: backpropagation through any operation that uses the safe Cholesky must produce finite gradients.
- The behavior should be the default and require no extra user configuration.

**Example behaviour that should work after the fix:**
Calling `sequential_gaussian_filter_sample` with the following 4×4 transition matrix (provided in a gist) for up to 1,000,000 steps should produce finite samples and finite gradients:

```python
trans_matrix = torch.tensor([...])  # ill-conditioned, see gist
init = Gaussian(zero, zero.expand(4), eye)
trans = matrix_and_gaussian_to_gaussian(trans_matrix, Gaussian(zero, zero.expand(4), eye)).expand((num_steps-1,))
x = sequential_gaussian_filter_sample(init, trans, (), noise)
# x must be finite, grads must be finite
```

No failures due to singular Cholesky are acceptable when the model is theoretically sound.

Add jitter to Cholesky factorization in Gaussian ops (#3151)


Implement `poutine.do` using node-splitting so that it models single-world counterfactual interventions. The new behavior must satisfy:

- **Intervention semantics**: For each intervened sample site, the model should use the provided intervention value as if it were hard-coded, but a fresh sample site (with the same name) must be introduced whose value does not propagate further. This preserves the original stochastic structure for downstream counterfactual reasoning.
- **Composition with `condition`**: `do` must compose freely with `condition` so that a single site can be both observed (factual) and intervened (counterfactual), enabling queries over potential outcomes.
- **Name mangling**: The original intervened node should be renamed (e.g., by appending `"__CF"`) to avoid collision with the newly introduced sample site, while the fresh sample site keeps the original name.
- **Warnings and errors**: If a site is intervened more than once, issue a `RuntimeWarning`. Interventions of non-numeric types (anything not a `Number` or `torch.Tensor`) must raise `NotImplementedError`.
- **Backward compatibility**: The old `do` behavior (hiding intervened sites via `block`) should be removed entirely; the new handler replaces it.

The change should be confined to the `poutine` module, introducing a new `DoMessenger` class and updating the `do` function. No changes to `ConditionMessenger` are required beyond removing the assertion that prevented overriding observed values, which is now handled correctly by stack overwrite behavior.

Rewrite poutine.do to represent single-world counterfactuals (#2157)


**Title:** ADVI guide ignores support transform Jacobian, yielding wrong parameter estimates for constrained latent variables

**Description**

When using ADVI (`advi_class` such as `ADVIDiagonalNormal` or `ADVIMultivariateNormal`) on a model that contains a latent variable with *constrained* support (e.g., `LogNormal`, `Beta`, `Gamma`, or any distribution where `site["fn"].support` is not `real`), the ELBO gradient is biased because the guide’s `log_prob` omits the `log_abs_det_jacobian` of the support transform.

Concretely, the current guide implementation applies `transform_to(support)` to the unconstrained parameter, then wraps the resulting value in a `Delta` distribution. That `Delta` alone only accounts for the point mass; it does **not** include the Jacobian correction that arises from mapping the unconstrained space to the constrained space. As a result, the learned variational parameters (e.g., `advi_loc` and `advi_cov`/`advi_scale`) do not recover the true generating parameters even for a trivial model like:

```python
def model():
    pyro.sample("x", dist.LogNormal(0.2, 0.7))
```

**Expected behaviour**

After running ADVI the learned variational posterior should match the true distribution’s parameters (within numerical tolerance). For the simple `LogNormal` model above, `advi_loc` should be close to 0.2 and the marginal diagonal variance should be close to 0.49.

**Suggested fix guidance**

The routine that constructs the guide’s `pyro.sample` call needs to produce a distribution whose `log_prob` on the sampled constrained value includes the support transform’s `log_abs_det_jacobian`. A natural way to achieve this is to **not** apply the transform eagerly to the value, but instead wrap the `Delta` distribution in a `TransformedDistribution`:

```
TransformedDistribution(Delta(unconstrained_value), transform_to(support))
```

This ensures the ELBO automatically incorporates the Jacobian term via the `TransformedDistribution.log_prob`.

Additionally, the current `Delta.log_prob` implementation performs aggressive broadcasting to the input `x` size, which can cause shape mismatches when `Delta` is used inside a `TransformedDistribution`. Simplify `Delta.log_prob` to compare `x` against `self.v` expanded to `self.shape()` (without broadcasting to `x.size()`). The expression `(x == v).type_as(x).log()` should suffice.

**Acceptance criteria**

- ADVI on a model with a constrained latent (e.g., `LogNormal`) recovers the true mean and variance within a few percent.
- No change to the signature or behaviour of the ADVI public interface.
- The guide continues to work unchanged for unconstrained (real-valued) latents.


Add a dedicated SVI inference path to `CompartmentalModel` alongside the existing MCMC path.
Rename the current `fit` method to `fit_mcmc` (keep its behaviour unchanged).
Implement a new `fit_svi` method that:

* Uses `AutoLowRankMultivariateNormal` as the variational guide and `ClippedAdam` as the optimizer.
* Accepts `num_steps`, `num_particles` (ELBO particles), `learning_rate`, `learning_rate_decay`, `betas`, `init_scale`, `guide_rank`, `haar`, `jit`, `log_every`, plus the usual heuristic options.
* Returns a list of per-step losses to allow convergence diagnostics.
* Works with the relaxed model, the Haar wavelet reparameterizer, and the internal `particle_plate` used for multiple posterior samples.
* Produces a `.samples` dictionary identical in structure to the one produced by `fit_mcmc`.

Fix shape handling in `_relaxed_model` and `_sequential_model` so that the particle dimension is correctly interleaved with time/region dimensions when the model is called inside a `particle_plate` (needed for SVI’s particle batching).

Relax the double-precision requirement so that it is only enforced when `num_quant_bins > 1` (i.e., only for MCMC quantized inference).

Update the example scripts `sir.py` and `regional.py` to support an `--svi` flag and the new SVI-related arguments; rename `--num-particles` to `--smc-particles` to avoid confusion with SVI’s `num_particles`. Update documentation to reflect that both inference backends are available.

Implement SVI inference for CompartmentalModel (#2529)


Fix `Categorical.sample` for numpy object-value spaces. When `_vs` is a numpy array of values, the returned sample should be a single object (not wrapped in a list) if the value space has only one row (i.e., `vs` is 1D or has batch size 1). For a true batch of value sets (more than one row), return a flat list of objects — no extra nesting. The current code always wraps the result in an extra list layer, which breaks downstream consumers like `Search._dist()` and makes trace sampling unusable. Adjust the logic to handle the single-row case and ensure batched output is a plain list. Also, remove any leftover `pdb.set_trace()` calls from the inference code.

fix to categorical sampling of objects, sampling tests now passing


We are experiencing several issues with the `pyro.infer.Predictive` class that hinder its usability, especially when exporting models. The class currently does not inherit from `torch.nn.Module`, which prevents it from being traced and serialized with `torch.jit`. Additionally, the semantics for `return_sites` when a guide is provided are inconsistent and missing guardrails, leading to confusing behaviors.

### Observed problems

1.  **Missing `nn.Module` base class** – `Predictive` cannot be used with `torch.jit.trace` or `torch.jit.trace_module`, making it impossible to save the predictive distribution as part of a serializable model.

2.  **No `forward` method** – The main callable is `get_samples`. This is incompatible with `nn.Module`’s `forward` interface and does not offer a JIT‑friendly signature that returns a `tuple` instead of a `dict` (a workaround required until PyTorch JIT supports dict outputs).

3.  **Undocumented special signal for `return_sites`** – When a guide is supplied, the internal logic uses a magic value (`return_sites=''`) to tell `_predictive` to return all sample sites. This is not exposed and conflicts with the type annotations.

4.  **No validation of conflicting arguments** – If a user provides both `posterior_samples` and a `guide`, the code silently misbehaves. This combination should be rejected because the guide already defines the posterior.

5.  **Ambiguous default behavior with guides** – It is unclear which sites are returned when a guide is present. A natural expectation is that when a guide is used, all latent sites from the guide should be sampled and returned (unless explicitly filtered), but the current default `return_sites=None` does not distinguish between “no filter specified” and “return nothing.”

### Expected behavior

- `Predictive` should be a `torch.nn.Module` so that it can be traced and serialized via `torch.jit.trace_module`. A `call` method that returns a `tuple` of sorted values should be provided as a workaround for JIT’s limitation on dict returns.
- The main entry point should be a `forward` method; `get_samples` may remain as a deprecated alias.
- When a `guide` is provided:
    - `posterior_samples` must be `None` or empty; passing both should raise a `ValueError`.
    - The default `return_sites` (an empty tuple) should cause all latent sites observed in the guide’s trace to be returned.
    - If a non‑empty collection is given for `return_sites`, only those specific sites should be returned (consistent with the no‑guide case).
- Without a guide, the existing behavior should remain unchanged, with the default returning all sample sites that are not already in `posterior_samples`. The internal signaling must be clean and not rely on string sentinels.

These changes will make `Predictive` usable in production pipelines that require JIT serialization and will provide clear, consistent interfaces for posterior predictive sampling.


```
Bug Report: Sampling from SpanningTree distribution raises TypeError when no initial edges are provided.

Description:
When calling `SpanningTree(edge_logits=...).sample()`, or equivalently the helper `pyro.distributions.spanning_tree.sample_tree(edge_logits, init_edges=None)`, the operation fails with a TypeError:

    sample_tree_approx() missing 1 required positional argument: 'edge_logits'

This occurs because `sample_tree` internally calls `sample_tree_approx(backend=backend)` – dropping the `edge_logits` argument – when `init_edges` is `None`. The intended behavior is to use `sample_tree_approx(edge_logits, backend=backend)` to generate an initial spanning tree and then apply the specified number of MCMC steps.

Acceptance Criteria:
- Calling `sample_tree(edge_logits, init_edges=None, ...)` no longer raises a TypeError.
- The function correctly uses the edge logit tensor to approximate an initial tree when no initial edges are supplied.
- The combination of approximate sampling followed by MCMC steps works as expected, allowing users to sample from a `SpanningTree` distribution without providing an explicit initial tree.
```

Minor fix to SpanningTree distribution (#1878)


The `rsample` method of `GaussianHMM` scales poorly with `obs_dim` when the observation distribution is `Normal(loc, scale).to_event(1)`. Sampling a moderate number of steps with an observation dimension of a few hundred already leads to tens of seconds of runtime and memory usage that can exceed available GPU memory.

**Observed behavior**:
Inside `GaussianHMM.rsample`, the code currently builds a full joint Gaussian over hidden and observed variables (by permuting and conditioning an `AffineNormal`), then draws a sample from that joint Gaussian. This involves a Cholesky factorization of a `(hidden_dim + obs_dim) × (hidden_dim + obs_dim)` precision matrix at each step, leading to `O(obs_dim³)` complexity and high peak memory.

**Expected behavior**:
Because the observation distribution is a diagonal `Normal`, the conditional distribution of observations given the hidden state is just an independent normal with no additional matrix algebra. The sampler should be able to directly sample observations in `O(obs_dim)` after the hidden states are drawn, without forming or factorizing the large joint precision matrix.

**Impact**:
Models with high-dimensional observation spaces (e.g., hundreds or thousands of independent sensor readings per time step) are unusable. Users are forced to reduce `obs_dim` drastically or avoid `GaussianHMM` entirely, even when the model is otherwise appropriate.

It would be helpful if `rsample` could detect when the observation distribution factors into independent normals and exploit that to avoid the expensive cubic operation, making sampling feasible for large observation dimensions.

Improve the performance of rsample of GaussianHMM with Normal obs (#2376)


Add `AVFMultivariateNormal`, an extension of `MultivariateNormal` that uses adaptive velocity fields as a control variate for lower-variance reparameterization gradients. The distribution must expose an additional trainable parameter `control_var` with shape `(2, L, D)` where `L` is a positive integer, validated at construction. The `rsample` method must implement custom forward/backward logic that adjusts the gradient of `scale_tril` using the control variate, while `control_var` itself receives gradients so it can be optimized jointly with `loc` and `scale_tril`. Also, fix `OMTMultivariateNormal` to pass `scale_tril` directly to the base `MultivariateNormal` constructor (instead of constructing a covariance matrix) and to raise `ValueError` on invalid input shapes rather than using `assert`.

add multivariate normal distribution with avf control variate (#1157)


**Title:** Add dynamic `.batch_shape()` and `.event_shape()` to distributions

**Context:**
Pyro distributions currently rely on a static `batch_size` argument passed to `__init__`, but the true batch shape is determined by the leading dimensions of the distribution’s parameter tensors. This creates ambiguity, makes manual construction error-prone, and prevents shape introspection needed for inference algorithms, shape-checking utilities, and user code.

**Requested behavior:**
1. Add a `.batch_shape(*args, **kwargs)` method to `Distribution` that returns a `torch.Size` of the batch (leftmost) dimensions derived from the parameters.
2. Add an `.event_shape(*args, **kwargs)` method that returns the event (rightmost) dimensions.
3. Add a `.shape(*args, **kwargs)` convenience method that returns `batch_shape + event_shape`, giving the sample tensor’s full shape.
4. Ensure:
   - `sample(*args, **kwargs)` returns a tensor of shape `self.shape(*args, **kwargs)`.
   - `batch_log_pdf(x, *args, **kwargs)` returns a tensor of shape `self.batch_shape(*args, **kwargs) + (1,)`.
   - The relationship `shape == batch_shape + event_shape` holds for all distributions.
5. Remove reliance on the static `batch_size` parameter where practical (or deprecate it). The new methods must compute shape purely from the input parameters, ignoring any stored `batch_size`.

**Scope:**
- Implement the abstract methods on `Distribution` and concrete implementations on `Categorical`, `DiagNormal`, `Dirichlet`, and any other core distributions needed for common usage.
- Update docstrings (especially in `Distribution`) to explain batch vs event axes and the shape guarantees.
- Ensure existing uses of `batch_size` do not break inference; the new methods must serve as the source of truth for shape queries.

**Edge cases / invariants:**
- For `Categorical` with `one_hot=True`, event shape is `(num_categories,)`; with `one_hot=False`, event shape is `(1,)`. Batch shape is the parameter shape minus the event dimensions.
- For `DiagNormal`, event shape is `(d,)` where `d` is the number of components. Batch shape is the parameter shape with the last dimension removed.
- For `Dirichlet`, event shape is `(d,)` similarly.
- Handle scalar parameters (zero batch dimensions) correctly (`.batch_shape()` returns `torch.Size([])`).
- The shape methods must accept the same optional arguments as `sample` / `batch_log_pdf` (e.g., `ps`, `one_hot` for `Categorical`, `mu`, `sigma` for `DiagNormal`). If arguments are omitted they should use default parameters stored at construction time, maintaining backward compatibility.

Add .batch_shape and .event_shape methods (#276)


We need a Delta distribution that can be used through the standard `RandomPrimitive` wrapper just like other distributions (e.g., Bernoulli, Beta), rather than being a stateful singleton instantiated directly. The current direct use of `Delta()` causes inconsistencies in tracing and reparameterization pipelines because the wrapper expects a class that is called with constructor arguments to build instances.

**Observed behavior that must change:**
- `delta` should be an instance of `RandomPrimitive(Delta)`, so constructing a distribution is done via `dist.delta(v)`, where `v` is the point value (a `Variable` or tensor). The underlying `Delta` class must accept `v` as its primary parameter at construction time and not as an optional instance attribute that mutates state.
- `sample()` should return the stored value `v` without extra arguments. The distribution must not require passing `v` again on each call to `sample` or to `log_pdf`.
- `batch_log_pdf(x)` must implement the proper point-mass log-probability: `0` for every element where `x == v` (within floating point tolerance), and `-inf` for any element not matching. The current approach of returning a large finite number (e.g., `999999`) is incorrect and breaks algorithms that rely on exact `-inf` as a sentinel for impossible values.
- The distribution should correctly expose `batch_shape` and `event_shape` so that `RandomPrimitive.shape` works and downstream inference algorithms can handle batching.
- The `support` method must return the variable representation of the stored value without requiring an extra argument.
- The `log_pdf` per-sample method does not need a custom override; the base `Distribution` class already derives it from `batch_log_pdf`. Any custom implementation that doesn't follow this convention should be removed.

**Constraints:**
- Backward compatibility is not required for the old singleton-style usage because it was broken with respect to the expected `RandomPrimitive` contract.
- The implementation must work with PyTorch `Variable`s and handle the case where the value `v` is a plain tensor (converting it to `Variable` if needed).
- The batch shape expansion logic should accommodate broadcasting when `x` and `v` sizes differ, to properly handle batched evaluation (e.g., multiple batch samples of points against a distribution with multiple possible values `v`).

RandomPrimitive wrapper and other fixes for Delta (#365)


```
The AIR example throws an error when run with the `--no-baseline` flag. To reproduce:

    python examples/air/main.py --num-steps=1 --no-baseline

Expected behavior: training should run without baseline-dependent computations, as the flag name implies.
Observed behavior: the script crashes with an `AttributeError` because a `None` value is used in place of a tensor.

The only requirement is that passing `--no-baseline` allows a full training run (any number of steps) without errors, while leaving the default baseline-enabled path unchanged. There is no need to change any other example or inference behavior.
```


Title: Enable composing EM, SVI, and multi-optimizer strategies with Newton and Adam

Background
We have building blocks for advanced inference: `Newton` for block-wise second-order optimization, `MixedMultiOptimizer` for applying different optimizers to different parameters, `TraceEnum_ELBO.differentiable_loss` for E-step computations, and `poutine.block` for hiding parameters from an outer optimizer.  However, composing them into reliable EM, nested EM-inside-SVI, and mixed-optimizer SVI loops is currently untested and may suffer from gradient-tracking issues in the parameter update machinery.

Request

1. **Make `MultiOptimizer.step` safe for gradient flow**
   - `MultiOptimizer.step` must copy updated parameter values into the parameter store **without** retaining autograd history that would incorrectly propagate gradients across iterations.
   - The current `value[...] = updated_values[name]` in-place assignment can leak gradients; replace it with a form that detaches the updated value before copying, e.g., `value.copy_(updated_values[name].detach())`.

2. **Ensure `Newton.get_step` returns differentiable results**
   - `Newton.get_step` currently detaches its output, preventing further differentiation (e.g., when used inside an outer SVI step).  Remove the `.detach()` so that the returned updated value remains differentiable when needed, while `MultiOptimizer.step` handles detachment for safe parameter storage.

3. **Provide worked examples of composition**
   - **Pure EM**: E-step via `TraceEnum_ELBO.differentiable_loss`, M-step via `Newton.step` on object locations.  Detach the parameter between EM iterations to break the autograd graph.
   - **EM inside SVI**: use `poutine.block` to hide M-step parameters from the outer SVI optimizer, run a few Newton steps on a detached copy of those parameters, then feed the updated (differentiable) values back into the parameter store so SVI can proceed on other parameters.
   - **SVI with mixed optimizers**: use `MixedMultiOptimizer` to combine Adam (e.g., for `noise_scale`) and Newton (for `objects_loc`) within the same SVI loop, using `poutine.trace(param_only=True)` to capture parameters and their unconstrained forms for the multi-optimizer.

   All examples should use a tracking model with object existence, data association, and `MarginalAssignment` for belief propagation, surfacing the relevant API patterns.

Acceptance criteria
- Existing single-optimizer uses of `MultiOptimizer` and `Newton` do not regress.
- The three compositional patterns can be run end-to-end, updating the correct parameters with the correct optimizers, and gradients flow correctly through the Newton M-step when required without accumulating across unrelated steps.

Add examples of composing inference algorithms (BP, EM, SVI, Adam, Newton) (#1247)


The `Trace` class used to record stochastic function execution in Pyro produces a memory leak during training loops. After repeated inference steps—and even after calling `gc.collect()`—the number of live `Trace` objects (and underlying `networkx.DiGraph` objects) grows without bound, eventually causing out-of-memory failures. The leak is reproducible in any workflow that creates and discards traces, including `poutine.trace`, `Trace.copy`, and SVI inference.

**Expected behavior**

- Repeated creation of `Trace` instances (e.g., via `poutine.trace(model)(data)` or `Trace.copy()`) must **not** cause a monotonic increase in the number of live `Trace` objects.
- After execution of a trace-producing call, when the caller no longer holds a reference to the trace, the trace and all its internal graph data must be eligible for garbage collection. A forced full GC should return the object count to a stable baseline.
- The same invariant must hold for `networkx.DiGraph` objects that back the trace; they must not leak independently.

**Scope and constraints**

- The fix must preserve the existing public API of `Trace`: all callers that access `nodes`, `edges`, `graph`, `in_degree`, `successors`, iteration, `__contains__`, `__len__`, `add_node`, `copy`, and the various stochastic/observation properties must continue to work without change.
- The graph ordering guarantee (via `collections.OrderedDict`) must be preserved.
- The `log_pdf`, `batch_log_pdf`, `compute_batch_log_pdf`, `compute_score_parts` methods must behave identically.
- The solution should not require changes to the workflow of external code that uses `Trace` objects.

Implement gross workaround to networkx memory leak


```
Fix `batch_log_pdf` shape and masking contract for `Bernoulli`, `DiagNormal`, and `Categorical`.

- When a `log_pdf_mask` is provided to `Bernoulli.batch_log_pdf` or `DiagNormal.batch_log_pdf`, the current implementation returns `torch.sum(log_pdf_mask * logsum, -1)` directly, skipping the normal sum-and-view logic. This breaks the output shape contract (it should end with `batch_shape + (1,)`).
  Instead, apply the mask by multiplying the per-element log values *before* the existing sum and view steps, letting them produce the correct final shape.

- `Categorical.batch_log_pdf` does not accept a `log_pdf_mask` at all. Add a `log_pdf_mask` parameter and, after computing the per-sample log probabilities, multiply by the mask (reshaped to match the batch shape) before returning.

- The mask is a tensor whose values scale the log probability: 0 to ignore a sample, 1 to keep, fractional values for KL annealing. It should be broadcast‑compatible with the per-sample log values (`Bernoulli`/`DiagNormal`) or with the final batch log-pdf tensor (`Categorical`). No early returns – the output must always go through the shape-normalization path used when the mask is absent.

- All other distributions remain unchanged (only these three previously exposed `log_pdf_mask` in their signatures).

- Ensure the change does not alter behaviour when `log_pdf_mask` is omitted, and that the masked output shape matches the `batch_shape + (1,)` invariant for every supported distribution.
```


We need three distributions—Binomial, Poisson, and LogNormal—fully available in `pyro.distributions` as thin wrappers around the corresponding `torch.distributions` implementations. This brings them in line with other distributions (e.g., Normal, Beta, …) and reduces custom code maintenance.

### Acceptance criteria

1. **Binomial**
   - Accessible as `pyro.distributions.Binomial` and via `pyro.distributions.binomial` (a `RandomPrimitive`).
   - Constructor signature: `Binomial(n, ps, ...)` where `n` is a positive integer (number of trials) and `ps` are `Variable`s holding success probabilities in [0,1].
   - Behaves as a discrete distribution compatible with the existing test harness.

2. **Poisson**
   - Accessible as `pyro.distributions.Poisson` and via `pyro.distributions.poisson` (a `RandomPrimitive`).
   - Constructor signature: `Poisson(lam, ...)` where `lam` is a positive rate (`Variable`).
   - Behaves as a discrete distribution.

3. **LogNormal**
   - **Replace** the custom implementation (`pyro.distributions.log_normal`) with a wrapper around `torch.distributions.LogNormal`.
   - The wrapper lives in `pyro/distributions/torch/log_normal.py`.
   - The old custom module must be removed.
   - Accessible as `pyro.distributions.LogNormal` and via `pyro.distributions.lognormal` (a `RandomPrimitive`).
   - Constructor signature `LogNormal(mu, sigma, ...)` unchanged.
   - **Preserves `analytic_mean()` and `analytic_var()`** methods that return the exact mean and variance without sampling, as in the old custom version.

4. **Public API**
   - All three classes are imported and exported in `pyro/distributions/__init__.py` alongside other wrapped distributions.
   - `RandomPrimitive` aliases (`binomial`, `poisson`, `lognormal`) are registered exactly as for other distributions.

5. **Documentation**
   - `docs/source/pyro.distributions.txt` is updated to include automodule entries for Binomial, Poisson, and LogNormal, pointing to their new `torch/` module paths.
   - The LogNormal entry uses `pyro.distributions.torch.log_normal`; Binomial and Poisson use `pyro.distributions.torch.binomial` and `pyro.distributions.torch.poisson` respectively.

### Non-goals
- No changes to distribution parameters semantics beyond using `torch.distributions` implementations.
- No changes to the testing infrastructure; the existing continuous/discrete test fixtures will verify correctness once the distributions are registered.

Reinstating Binomial, Poisson and LogNormal from torch.distributions (#720)


Queue poutine is broken. We have a commented-out test suite for it in `tests/test_poutines.py` (search for `QueuePoutineTests`). The symptom is that `poutine.queue` cannot build a correct trace for enumeration, and basic usage fails in several ways.

Diagnosis from a quick look suggests:
- The `Trace` class’s `add_*` methods (`add_sample`, `add_observe`, `add_map_data`, `add_param`, `add_args`, `add_return`) do not return `self`, which prevents method chaining. This breaks trace accumulation inside `QueuePoutine` and likely other callers that expect a fluent API.
- `pyro.infer.Trace()` is not properly importable/accessible via the public module path (for example, `pyro.infer.Trace()` fails). The queue poutine needs to instantiate a new `Trace` from the `pyro.infer` namespace.
- There are Python 3 incompatibilities in `QueuePoutine`: `xrange` used where `range` is needed, and `max_tries` is set as a float but later used in a `range` call, which would fail.
- Several files still contain `import pdb` left over from debugging.

The request:
1. Make `Trace` fluent so that each `add_*` method returns the trace instance. This should be enough to let `QueuePoutine` and similar poutines chain trace construction correctly.
2. Ensure `pyro.infer.Trace` is exported properly so that `pyro.infer.Trace()` works, and adjust `QueuePoutine`’s internal trace creation to use that public path.
3. Fix the Python 3 issues in `QueuePoutine` (xrange, max_tries type).
4. Clean up the dead `import pdb` statements so that the queue poutine test module doesn’t reintroduce them.
5. After these fixes, the commented-out `QueuePoutineTests` in `tests/test_poutines.py` should become passable. Those tests cover: running a single trace through a queue poutine, enumerating all discrete paths (a toy HMM with 2^3 possibilities), triggering a `ValueError` when `max_tries` is exceeded, and queue exhaustion. Enable those tests and make sure they pass.

No API changes beyond what’s needed to fix the broken behavior; the public signatures of `poutine.queue`, `Trace.add_*`, and `pyro.infer` should remain backward-compatible in intent.


We need a circular von Mises distribution in Pyro.
Implement a `VonMises` class that inherits from `TorchDistribution`.

**Parameters**
- `loc` (angle in radians, unconstrained)
- `concentration` (positive)

**Scope**
Only `log_prob(value)` is required for now — sampling, `mean`, `variance`, etc. can be omitted.

**Normalization**
The log‑probability must use a reliable numerical approximation of \(\log I_0(\kappa)\) (the log of the modified Bessel function of order 0) to keep the density correctly normalised across a wide range of concentrations (from very small to large values).

**Registration**
The new distribution should be:
- importable from `pyro.distributions` (i.e. added to `pyro/distributions/__init__.py`)
- documented in the Sphinx documentation (add an entry to `docs/source/distributions.rst` under “Primitive Distributions”).

**Constraints**
Use `arg_constraints = {'loc': constraints.real, 'concentration': constraints.positive}` and `support = constraints.real`.

This will let users express circular angle likelihoods directly in Pyro models and reuse the existing distribution testing infrastructure for validation.

Add circular von Mises distribution (only .log_prob() is implemented) (#1038)


Fix JIT compatibility across examples and internal utilities so that more models can run with the `--jit` flag without errors.

### Problem areas

- **LDA example** (`examples/lda.py`) fails under JIT tracing because its model uses an `assert data.shape` statement and its guide uses `scatter_add_`, neither of which are traceable. The example must support `--jit` without crashing; when tracing is active, it must use alternative code paths that are JIT-safe. Also, the default `--num-words-per-doc` produces a shape mismatch in the guide’s histogram construction under JIT; increase the default from 32 to 64 to avoid this mismatch.

- **`scale_and_mask`** in `pyro/distributions/util.py` still carries a workaround for PyTorch issue [#12230](https://github.com/pytorch/pytorch/issues/12230), which has been resolved. The function conditionally uses `tensor[~mask] = 0` when traced instead of `masked_fill_`. Remove the workaround – call `masked_fill_` unconditionally and drop the extra `.contiguous()` call.

- **`JitTraceEnum_ELBO`** in `pyro/infer/traceenum_elbo.py` uses `elbo += ...` (in-place addition) inside a JIT-compiled function, which Torch JIT does not support. Replace it with `elbo = elbo + ...`.

- **`_Marginal.diagnostics`** in `pyro/infer/mcmc/mcmc.py` crashes when `effective_sample_size` raises `NotImplementedError`. Instead, catch the exception and set `n_eff` to `NaN` so diagnostics remain available for other quantities.

- **Example tests** (`tests/test_examples.py`): `eight_schools/svi.py --jit` is currently marked `xfail`; after the fixes it should run normally. `eight_schools/mcmc.py --jit` is too slow to run and should be skipped instead of being marked as an expected failure.

More jit fixes, e.g. baseball and eight_schools/svi (#1598)


Implement a `SineSkewed` distribution that wraps any existing torus base distribution (support on a product of circles) to break symmetry via sine skewing as described in Ameijeiras-Alonso & Ley (2019). The new distribution must:

- Accept a `base_dist` (any `TorchDistribution` with event shape `(d,)` and range `[-π, π)^d`, e.g. `Uniform(-π,π)`, `VonMises`, `SineBivariateVonMises`, `ProjectedNormal`).
- Accept a `skewness` tensor of shape `batch_shape + (d,)` with each element in `[-1, 1]`. The constraint `abs(skewness).sum(-1) <= 1` should be enforced (a warning is acceptable if violated; no automatic renormalization).
- `log_prob(value)` should add the skewing term: `log(1 + (skewness * sin(value - base_dist.mean mod 2π)).sum(-1))` to `base_dist.log_prob(value)`.
- `sample` should use the rejection-free algorithm from the paper: draw `y ~ base_dist`, draw `u ~ Uniform(0,1)`, compute `mask = u ≤ 0.5 + 0.5 * (skewness * sin(y - base_dist.mean mod 2π)).sum(-1)`, and return `y` where `mask` is true, else `-y + 2*base_dist.mean` wrapped back to `[-π, π)`.
- Support `expand` to new batch shapes.
- Document that the distribution is not reparameterized, so it will be slow for latent variable inference in SVI, but can be used as a likelihood and is suitable for HMC/NUTS.

The class should be registered in `pyro.distributions` and exposed in the documentation alongside other torus distributions.

Sine Skewed Toridial distribution (#2826)


The `pyro.infer.MH` sampler produces biased posterior samples when a guide is provided. For a simple Normal-Normal conjugate model where the guide proposes from the prior, the estimated posterior mean and standard deviation are far from the true values, even after thousands of samples. Inspecting the Metropolis-Hastings acceptance logic in `pyro/infer/mh.py`, the log acceptance ratio appears to be computed as:

```
logr = new_model_trace.log_pdf() + new_guide_trace.log_pdf() - old_model_trace.log_pdf() + old_guide_trace.log_pdf()
```

For a proposal \(q(z'|z)\), the correct log acceptance ratio should be:

\[
\log \frac{p(x,z')}{p(x,z)} + \log \frac{q(z|z')}{q(z'|z)}
\]

i.e., `log p(x,z') + log q(z|z') - log p(x,z) - log q(z'|z)`. The current code adds the old guide log-pdf instead of subtracting it, which corrupts the acceptance decision and pulls the chain away from the true posterior.

Additionally, the sample counter `t` is only incremented inside the `if t <= self.burn or …` condition that appends to the trace list. This means accepted samples that do not satisfy the burn-in / thinning condition never increment `t`, so the total count of accepted steps is undercounted. This reduces the effective number of collected samples beyond the intended `samples` and `lag` settings, further degrading estimates.

**Expected behavior:**
1. The log acceptance ratio should correctly implement the MH formula using subtraction of the old guide log-pdf.
2. The counter `t` should be incremented once per accepted proposal, independently of whether the sample is stored in the output trace list, so that burn-in and thinning behave as intended.

**Impact:**
Users relying on `MH` with a guide to estimate posterior moments get systematically wrong results, even for textbook conjugate models. The bug makes the sampler unfit for its advertised purpose.

fixed mh acceptance ratio and changed mh test tolerances


Add the half-Cauchy distribution (location `mu`, scale `gamma`) to `pyro.distributions`. It should follow the existing Cauchy implementation conventions: support `mu=None` for a stateful distribution, provide `sample` via `scipy.stats.halfcauchy.rvs` and `batch_log_pdf` matching the density `2 / (π γ (1 + ((x-μ)/γ)²))` for `x > μ`. `analytic_mean` and `analytic_var` must raise `ValueError` (undefined). Wire it into `pyro/distributions/__init__.py` as `HalfCauchy` with a module-level alias `halfcauchy`.


**Title:** Support `jit_compile=True` for HMC kernels when the model uses `pyro.param` tensors with `requires_grad=True`

**Description**
When using HMC-based kernels (e.g., `NUTS`) with `jit_compile=True`, the potential energy function is traced via `torch.jit.trace`. If the model contains global parameters created with `pyro.param` that retain their default `requires_grad=True` flag, the tracing step currently fails or produces an unusable compiled function. This makes it impossible to benefit from JIT compilation in models that simply declare a differentiable parameter used in sampling statements.

**Acceptance Criteria**
1. Running `MCMC` with `NUTS(model, jit_compile=True)` (and any other HMC-derived kernel) succeeds for models that contain `pyro.param` tensors with `requires_grad=True`.
2. After the MCMC run (or after the potential function is compiled), the `requires_grad` flags of the parameter store’s tensors remain unchanged from their original state.
3. The JIT-compiled potential function produces identical evaluated values (up to floating point tolerance) as the non-JIT-compiled version for the same inputs.
4. The solution must not alter how models without `pyro.param` or with `requires_grad=False` tensors behave—those should still compile and run as before.
5. The compilation step should not leak gradient information into the traced graph; the resulting `ScriptModule` should be independent of any surrounding autograd context.

**Current Behavior (pre-fix)**
Invoking `jit_compile=True` on a model that declares `pyro.param("x", torch.tensor(1.))` raises a `RuntimeError` from `torch.jit.trace` because leaf tensors that require grad are not directly traceable. Users must manually set `requires_grad_(False)` on all param tensors before constructing the kernel, which is error-prone and unfriendly.

**Additional Constraints**
- The fix must be thread-safe and work with multiple chains (`num_chains > 1`).
- Temporary modification of `requires_grad` flags must handle cases where some params require grad and others do not.
- Performance of the trace should not degrade significantly; the fix may temporarily disable gradients only for the duration of the `torch.jit.trace` call.

Allow jit compile models with params in HMC (#2355)


The OED A/B test example (`examples/contrib/oed/ab_test.py`) is flaky when run with small numbers of VI and BO steps (e.g., 10 VI steps, 2 BO steps). Occasionally the example crashes with a NaN error during the Bayesian optimization stage. This makes the example unreliable in automated testing and for quick experimentation.

**Observed behavior**
When the Bayesian optimizer proposes a design candidate, the objective function sometimes returns a NaN value. That NaN then propagates into the optimization state and causes a hard failure later in the run. The outcome depends on the random seed, and with the hardcoded seed the example may repeatedly hit the problem or rarely pass.

**Expected behavior**
- The Bayesian optimization loop should gracefully handle candidates that produce NaN objective values (e.g., by discarding them and trying another candidate), rather than crashing or producing incorrect results.
- The example should be stable enough to run without random failures in continuous integration.
- For reproducibility and testability, the example should accept a configurable random seed instead of a fixed constant.

**Additional context**
The flaky behavior is currently worked around by marking the example as an expected failure in tests. A fix in the optimizer utility used by several OED examples (`gp_bayes_opt.py`) would benefit all examples that rely on it.

Fix the flakiness of oed example (#1645)


**Subject:** Missing `GammaPoisson` compound conjugate distribution

**Problem**
The `pyro.distributions.conjugate` module currently provides a `BetaBinomial` compound distribution, which integrates a Beta prior over a Binomial success probability. There is no analogous compound distribution for a Gamma prior over a Poisson rate. This forces users who want to model overdispersed counts with a Gamma–Poisson hierarchy to manually compute log-probabilities or to work with the less convenient `NegativeBinomial` parametrization.

**Expected behavior**
A new `GammaPoisson` distribution class should be added, following the same pattern as `BetaBinomial`:
- Parameters: `concentration` (shape/α) and `rate` (rate/β) matching the Gamma distribution.
- The distribution should represent the marginal of a Poisson distribution whose rate is drawn from `Gamma(concentration, rate)`.
- It must support the standard distribution interface: `sample`, `log_prob`, `mean`, `variance`, `expand`, etc.
- It should be a discrete distribution with support on the non-negative integers.
- A documentation note should clarify that this is an alternative parametrization of the `NegativeBinomial` with `total_count = concentration` and `probs = 1 / (1 + rate)` (or equivalently `logits = -log(rate)`), so users can choose whichever parameterisation is most natural.

**Safety / constraints**
- Both `concentration` and `rate` must be positive (`constraints.positive`).
- `log_prob` must handle arbitrary batch shapes and validate arguments correctly.
- The implementation should reuse or wrap an internal `Gamma` instance for sampling and to derive the parameters, as `BetaBinomial` does for `Beta`.
- The `expand` method must correctly propagate the internal `Gamma` distribution.

**Impact**
Without this distribution, models that naturally use a Gamma prior on Poisson rates must resort to manual implementations or less intuitive parametrizations, creating friction and potential for error. Adding `GammaPoisson` would mirror the existing conjugate pair utilities and improve the library’s completeness.

Adding Gamma-Poisson distribution (#1713)


**Title:** Add `differentiable_loss()` to `TraceEnum_ELBO`

**Context:**
`TraceEnum_ELBO` currently exposes `.loss()` (returns a plain `float`) and `.loss_and_grads()` (populates `.grad` fields via autograd), but there is no way to obtain the ELBO as a raw differentiable `torch.Tensor`. This blocks advanced use-cases such as higher-order differentiation, custom loss re-weighting, or combining the ELBO with other differentiable penalties before calling `.backward()`.

**Request:**
Implement a `differentiable_loss(model, guide, *args, **kwargs)` method on `TraceEnum_ELBO` that:

- Returns a `torch.Tensor` representing the (estimated negative) ELBO, averaged over `num_particles`.
- The returned tensor must require gradient and carry a valid computation graph (i.e., be the result of differentiable operations over model/guide parameters).
- If the ELBO is identically zero for the given model/guide pair, the method should raise a `ValueError` indicating the ELBO is not differentiable.
- Otherwise, the tensor should be infinitely differentiable as long as the underlying distributions support it.

**Implementation note:**
The existing `_compute_dice_elbo` helper already computes a per-particle differentiable ELBO; the new method should reuse it in the same way `loss_and_grads` does, but return the summed-and-averaged tensor instead of calling `.backward()`. The method must also replicate the NaN-checking that `loss()` already performs.

Implement TraceEnum_ELBO.differentiable_loss() method (#1243)


I am using Pyro’s discrete enumeration to train Hidden Markov Models, following the patterns in `examples/hmm.py`. Models like `model_5` in that example use a neural network to emit observation logits and rely on `TraceEnum_ELBO` to marginalize out the latent state sequence. However, the current approach unrolls the time dimension into a loop of sample statements, one per time step. This leads to serial execution with linear complexity in the number of time steps, making training prohibitively slow for longer sequences.

I would like a new distribution, `DiscreteHMM`, that computes the log-probability of an HMM in a *parallel* manner, reducing the critical path to O(log T) using the temporal parallelization technique from Särkkä & García-Fernández (2019). This distribution should:

1. **Interface**:
   - Accept `initial_logits` (state_dim,), `transition_logits` (num_steps, state_dim, state_dim) and an arbitrary `observation_dist` whose `.batch_shape[-1]` equals `state_dim`.
   - Support heterogeneous transition and observation distributions across time. For the fully homogeneous case, the `event_shape` should become `(1,) + observation_dist.event_shape`, allowing `log_prob` to accept data of any length.

2. **Event shape**:
   - Include the time dimension on the left: `event_shape = (num_steps,) + observation_dist.event_shape`.
   - When time-varying parameters are broadcastable (homogeneous), the single-step `event_shape = (1, …)` should be compatible with variable-length data.

3. **Masking**:
   - Be compatible with `MaskedDistribution` (i.e., allow `obs_dist.mask(...)`) so that sequences of unequal length can be handled in a batch without manual truncation.

4. **Numerical stability**:
   - Implement log-probability entirely in log-space using stable `logsumexp`–style operations, avoiding underflow or overflow on long chains.

5. **Integration**:
   - Be importable from `pyro.distributions` and documented in the distributions Sphinx page.
   - Work seamlessly inside a `pyro.sample` statement with `TraceEnum_ELBO` (or `JitTraceEnum_ELBO`), with no unexpected enumeration warnings (the distribution itself handles the marginalization over latent states).

**Expected impact**:
A new model (`model_7`) in `examples/hmm.py` that replaces the time-step loop with a single `pyro.sample("y", DiscreteHMM(...), obs=y)` should achieve >10× speedup during training compared to the equivalent sequential model (`model_5`). This will make discrete-state HMMs with long observation sequences much more practical, and the implementation can serve as a building block for other dynamic Bayesian networks.

Implement parallelized DiscreteHMM distribution (#1958)


The `Coregionalize` kernel is currently defined as \(k(x,z) = x^T W W^T z\), which is purely low-rank. In multi-output Gaussian process models this is often insufficient; a diagonal term \(D = \operatorname{diag}(d)\) is needed to get a full-rank kernel of the form \(k(x,z) = x^T (W W^T + D) z\). This diagonal component captures per-output variance that cannot be explained by the low-rank structure, improves numerical stability, and aligns with standard coregionalization formulations.

**Observed problem**
The class only accepts a `components` matrix (or a `rank` to randomly initialize one) and entirely lacks a diagonal parameter. As a result, the kernel is rank-deficient and cannot represent independent output-specific scales or account for output-level noise in a unified way.

**Expected behavior**
The constructor should also accept an optional `diagonal` argument – a 1‑D tensor of length `input_dim` – representing the diagonal entries added to the weight matrix. The kernel output should become
```
X @ (W @ W.T + diag(d)) @ Z.T
```
for the full covariance and similarly for the diagonal (`diag=True`) case.

**Additional requirements**
- The `diagonal` vector must be constrained to positive values.
- When `diagonal` is not specified, a sensible default should be used. A natural choice is to initialise it to a constant vector of 0.5, so that when `components` is randomly initialised with an expected outer product of 0.5 ⋅ I, the whole matrix \(W W^T + D\) has expected value the identity.
- The `forward` method needs to incorporate the diagonal term both in the cross-covariance case (`X, Z`) and in the `diag=True` variant, while preserving the existing logic for slicing active dimensions.

**Impact**
Without this enhancement, users have to manually add a separate diagonal (e.g., via a `WhiteNoise` kernel plus a `Sum` kernel), which is both cumbersome and not equivalent to a principled coregionalization kernel with a diagonal component. Adding the `diagonal` parameter directly to `Coregionalize` would make the kernel fully expressive for multi-output models and improve out‑of‑the‑box functionality.

Add diagonal component to coregionalization kernel (#1116)


We need a unified factory for the number of new infections that automatically chooses the right Pyro distribution (Poisson, Negative‑Binomial, Binomial, or Beta‑Binomial) based on whether the population is finite and whether superspreader overdispersion is requested.  Currently `SimpleSIRModel` and `SimpleSEIRModel` compute the infection probability manually:

```python
rate_s = -R0 / (tau * self.population)
prob_s = -(rate_s * state["I"]).expm1()
S2E = pyro.sample("S2E_{}".format(t), dist.Binomial(state["S"], prob_s))
```

This has two problems:

1. It only supports a standard Binomial (or implicitly Poisson when population is large), with no way to add the overdispersion needed for superspreading individuals (e.g. Negative‑Binomial or Beta‑Binomial).
2. The manual probability calculation is error‑prone, especially with clamping, and doesn’t correctly handle the large‑population Poisson limit in a uniform way.

The same pattern is duplicated in the forward and backward transitions of both models.

We would like a helper function `infection_dist` in `pyro.contrib.epidemiology.distributions` that encapsulates the four possible combinations:

- If the population is infinite and there is no overdispersion (`concentration = ∞`), return `dist.Poisson(individual_rate * num_infectious)`.
- If the population is infinite but overdispersion is requested (`concentration < ∞`), return `dist.NegativeBinomial(total_count = concentration * num_infectious, logits = log(individual_rate / concentration))`.
- If the population is finite and there is no overdispersion, return `dist.ExtendedBinomial(num_susceptible, 1 - (1 - individual_rate / population) ** num_infectious)` with appropriate clamping to avoid values ≥ 1.
- If the population is finite and overdispersion is requested, return `dist.ExtendedBetaBinomial` with parameters derived from `concentration`, `num_infectious`, and the finite‑population infection probability.

All edge cases (e.g., `p` very close to 1, `rate / N` near the upper limit) must be handled safely using log‑space and clamping.  The factory should expose a simple signature:

```python
infection_dist(*, individual_rate, num_infectious,
               num_susceptible=math.inf, population=math.inf,
               concentration=math.inf)
```

This will allow models to switch between epidemiological assumptions simply by passing or omitting `population` and `concentration`.  The existing `SimpleSIRModel` and `SimpleSEIRModel` should be updated to use this factory in both `transition_fwd` and `transition_bwd` instead of computing `prob_s` by hand.  This reduces code duplication and makes it trivial to experiment with overdispersed dynamics or finite‑size corrections.

Add a factory for Poisson / NegativeBinomial / Binomial / BetaBinomial (#2450)


We’ve observed that `Predictive.get_samples()` returns samples with incorrect tensor shapes when the model uses `pyro.plate`.
Specifically, the resulting sample tensors only carry the leading `num_samples` dimension, but are missing the batch dimensions that correspond to plates in the model.
For example, a site inside a plate of size 5 should have shape `(num_samples, 5, …)`, but instead appears with shape `(num_samples, …)`.
This happens both with `parallel=True` and `parallel=False`, though the exact malformed shape may differ between the two modes.

Expected behavior:
The shapes produced by `Predictive.get_samples()` should exactly match the shapes obtained from manually vectorizing the model under a `pyro.plate("_vectorize", num_samples, dim=…)` and replaying with a guide trace.
For any return site, the output tensor must reflect the full plate nesting that was present in the model, so that the result is usable for downstream computation without manual reshaping.

The issue makes it impossible to reliably use `Predictive` with plate-structured models, as the returned tensors cannot be broadcast or reduced consistently with the original model’s geometry.
Please ensure that when computing the output shapes, the plate nesting depth is taken into account, and that the sequential and parallel code paths produce congruent results.


The `ProjectedNormal.log_prob()` method can return NaN for numerically challenging but valid inputs, particularly when the concentration vector has large positive entries (e.g., strength ≫ 100) or when the evaluated point is nearly aligned with the concentration direction. The issue persists in both float32 and float64 precision, although it is easier to trigger with single-precision floats.

**Observed behavior**
- Calling `log_prob` with large concentration parameters (e.g., `[100, 200, 300]`) on points close to the mode sometimes yields `NaN`.
- Backpropagating through the log-probability also produces `NaN` gradients, which breaks gradient-based inference methods (SVI, HMC, etc.).

**Expected behavior**
`log_prob` and its gradient should be finite and well-behaved for any finite concentration vector and any value on the sphere, with no NaN or `-inf` for legitimate parameter settings.

**Likely root cause**
The current implementation in `_log_prob_2`, `_log_prob_3`, and `_log_prob_4` computes terms like `(1 + t * erf(t / sqrt(2))) / 2` where `t` can be large and negative, causing catastrophic cancellation in `1 + erf(...)`. Additionally, the final argument to `log` can become zero or slightly negative due to round-off, leading to `-inf` or `NaN`.

**Desired fix**
- Rewrite the relevant formulas using the complementary error function (`erfc`) to avoid the cancellation, replacing `(1 + erf(x))/2` with `erfc(-x)/2`.
- Clamp the argument to `log` to a small positive value (e.g., maximum of the expression and a tiny epsilon) to guard against floating-point underflow/rounding below zero.

This should make `ProjectedNormal.log_prob()` safe to use throughout the entire parameter range, including strong concentrations, and restore proper gradients for downstream inference.

Numerically stabilize ProjectedNormal.log_prob() via erfc (#3071)


**Bug: Custom subsample argument to iarange/irange causes error or ignores the provided indices**

## Description
When a user provides a custom list of indices via the `subsample` argument of `pyro.iarange` or `pyro.irange`, the code fails to iterate over that list correctly. This prevents deterministic batching while still using the cleaner iarange/irange context-manager API.

## Steps to reproduce
```python
import pyro

def model():
    with pyro.iarange("batch", 20, subsample=[1, 3, 5, 7]) as batch:
        result = batch
    return result

model()
```
or, using `irange`:
```python
def model_irange():
    result = []
    for i in pyro.irange("batch", 20, subsample=[1, 3, 5, 7]):
        result.append(i)
    return result
```

## Expected behavior
- `iarange` should yield the provided list as `batch` without performing any random sampling. The user should be able to use `batch` directly as a tensor/list of indices.
- `irange` should iterate over the elements of the provided list, yielding each index one by one.
- `poutine.trace` should work correctly with such a model (the traced model should return the same custom subsample list).

## Actual behavior
- `iarange` may call `sample(..., obs=subsample)` and thus potentially involve the sampling mechanism even when a fixed subsample is supplied.
- `irange` eventually tries to access the `.data` attribute on `batch`, which fails with an `AttributeError` when `batch` is a plain Python list.

## Environment
Pyro main branch, any recent PyTorch version.

Fix custom subsampling in iarange, irange


Fix `iarange` and `irange` when `subsample_size == size`. Currently, when the subsample covers the entire dataset, the index set is returned directly without touching the plate stack, which means:
- No scaling is applied, so nested plate log-prob corrections are missing.
- Sampled nodes inside the plate lack the `map_data_stack` attribute.
- Replay of a model that uses a full-size `iarange` fails or produces wrong results because the plate is invisible to the poutine trace.

The intended behavior:
- Even when `subsample_size == size`, the generated indices (deterministically `range(size)`) must flow through the `_PYRO_STACK` logic.
- Scaling and `map_data_stack` must be set on nodes exactly as they would be for any other subsample size.
- Replay should work: a replayed model should yield the identical deterministic indices, while a fresh trace may differ only when the size is smaller than the total.

The fix should apply to `iarange` directly; `irange` should inherit the correction automatically if it delegates to `iarange`.

Fix bug in iarange,irange when subsample_size == size


We need an example implementation of the sparse gamma deep exponential family (DEF) model from Ranganath et al., applied to the Olivetti faces dataset. The example should live in the `examples/` directory and be a self-contained script that:

- Loads the Olivetti faces training data. If the file `faces_training.csv` is missing, download it automatically using the `wget` library.
- Defines a three-layer Gamma DEF with Poisson observation likelihood, along with a mean-field Gamma guide, using Pyro’s `SVI` and `Trace_ELBO`.
- Accepts a `--num-epochs` command-line argument (default 1001) to control training length.
- Prints the training ELBO every 20 epochs after epoch 0.
- Clips variational parameters after each gradient step to avoid vanishing means and ensure numerical stability.
- Is runnable with a single epoch (e.g., `--num-epochs=1`) without errors, for smoke-testing purposes.

Additionally, since the script depends on `wget` for automatic data fetching, add `wget` to the project’s extra requirements (the `EXTRAS_REQUIRE` list) so that the example environment includes the necessary dependency.

add sparse gamma deep exponential family example (#1145)


**Feature request: Allow HMC and NUTS kernels to accept a user-provided potential function**

As a user, I sometimes need to draw MCMC samples from an unnormalized distribution that is not defined as a Pyro model (e.g., when the log-density is computed by an external library, or is expressed in closed form that doesn’t map naturally to Pyro primitives). Currently, HMC and NUTS require a Pyro model and derive the potential energy by tracing it. I want to bypass the model entirely and supply my own `potential_fn`.

**Acceptance criteria (observable behavior):**

1. **Providing `potential_fn` to kernel constructors**
   - Both `HMC` and `NUTS` accept a new optional argument `potential_fn`.
   - If `potential_fn` is given, it is used as the potential energy function; the `model` argument becomes optional (can be `None`).
   - If both `model` and `potential_fn` are supplied, `potential_fn` takes precedence for energy calculations. (Passing only `model` continues to work exactly as before.)

2. **Interface of `potential_fn`**
   - `potential_fn` is a callable that receives a dict mapping site names (strings) to torch tensors in the *unconstrained* space (real support), and returns a scalar torch tensor representing the potential energy.
   - The integrator uses this function directly (including automatic gradient computation) in place of the old model‑trace‑based evaluation.

3. **MCMC workflow without a model**
   - When a kernel is created with `potential_fn` and `model=None`, the `MCMC` runner must handle the absence of a prototype model trace.
   - During sampling, the `MCMC` runner should yield the raw unconstrained parameter dict directly (instead of attempting to convert to a constrained trace via `_trace_wrap`), since there is no model to constrain against.
   - This means users can directly consume the samples from `mcmc_run.exec_traces` as dicts of unconstrained tensors.

4. **Initial parameters**
   - When using only `potential_fn`, the kernel does not automatically generate initial parameters. The user must set `kernel.initial_params` manually to a dict of unconstrained tensors before calling `MCMC.run()`.
   - The kernel should accept and use these initial parameters for warmup and sampling without requiring a model trace to find valid starting points.

5. **Feature parity for adaptation and integration**
   - Step size adaptation, mass matrix adaptation (diagonal or dense), and the NUTS tree‑building algorithm all work identically regardless of whether the potential comes from a model or from the user‑provided function.
   - The JIT compilation option (`jit_compile`) should also be usable when only a `potential_fn` is provided, speeding up evaluation if desired.

6. **No breakage of existing model‑based workflows**
   - Any existing code that creates `HMC` or `NUTS` with a positional (or keyword) `model` argument and no `potential_fn` continues to work with no change in behavior.

**Constraints / edge cases:**
- If a `potential_fn` is provided that returns `NaN` or `Inf`, the kernel should behave as it would for a degenerate model (e.g., reject the proposal, attempt to find a valid initial state during warmup, etc.), following the same numerical safety rules that already exist.
- The `transforms` argument to the kernel is still accepted but may be deprecated/ignored when a model is absent; in any case, `potential_fn` operates on user‑supplied unconstrained values, so no automatic transform is applied.
- The `max_plate_nesting` argument is similarly irrelevant when there is no model, and should be gracefully ignored.

Revise HMC to support general potential_fn (#1845)


**Title:** Drop string-based `loss` parameter from `SVI` and require explicit `ELBO` instances

**Summary**
Currently, `SVI.__init__` accepts `loss="ELBO"` as a convenience, which internally dispatches through `ELBO.make()` and passes on extra keyword arguments. We want to simplify the API and remove the string path entirely. Callers should construct and pass an `ELBO` subclass instance directly (e.g., `Trace_ELBO()`).

**What needs to change**

1. **`SVI` (`pyro.infer.SVI`):**
   - Remove the branch that handles `isinstance(loss, str)`.
   - Accept only `ELBO` instances; raise a clear `TypeError` for anything else.
   - Stop accepting extra `**kwargs` that were previously forwarded to the factory (`trace_graph`, `enum_discrete`, etc.).

2. **`ELBO` (`pyro.infer.ELBO`):**
   - Remove the `make()` static method. It is no longer needed once `SVI` stops using it.

3. **Internal callers:**
   - `pyro.contrib.autoguide.advi` docstring example: replace `'ELBO'` with `Trace_ELBO()`.
   - `pyro.contrib.gp.models.GPModel.optimize`: currently hard-codes `loss="ELBO"`. Update the signature to accept an optional `loss` parameter (defaulting to `Trace_ELBO()`). Also handle `optimizer=None` by defaulting to `Adam({})`, matching the existing behavior but without relying on string-based loss.
   - Any docstring examples in `GPModel` or related files that show `SVI(... , loss="ELBO")` must be updated.

4. **Downstream compatibility:**
   - Users who were passing string `"ELBO"` will now see an error telling them to pass an `ELBO` instance. That is intentional, but the error message should be helpful.

**Edge cases / constraints**
- The `SVI` constructor should not silently accept old keyword arguments that were used for ELBO selection; after this change they become invalid.
- The `GPModel.optimize()` convenience must remain backward-compatible in terms of default behavior (Trace_ELBO equivalent).
- No other loss types are supported; only `ELBO` subclass instances are valid.

**Acceptance criteria**
- Passing `loss="ELBO"` to `SVI` raises a `TypeError`.
- `ELBO.make()` is removed.
- All internal references to `loss="ELBO"` (including docstring examples) are replaced with appropriate `ELBO()` instances.
- `GPModel.optimize()` works when called with no arguments (using default `Adam` and `Trace_ELBO`), and can be customized by passing explicit `loss` and `optimizer` objects.

Drop support for string value of loss arg to SVI (#1054)


**Title:** Missing `filter` method for HMM distributions and `validate_args` not preserved in `expand`

**Description:**
The `DiscreteHMM` and `GaussianHMM` distributions lack a method to compute the posterior distribution over the final hidden state given a sequence of observations. This is needed to perform forecasting: after observing a sequence, one should be able to obtain a distribution over the latest latent state and then treat it as the initial state distribution for a new HMM over future time steps. Without this, users must manually reimplement the filtering logic, which is error‑prone and breaks the encapsulation of the HMM classes.

**Expected behavior:**
- Add a `filter(self, value)` method to both `DiscreteHMM` and `GaussianHMM`.
- For `DiscreteHMM`, the method should return a `Categorical` distribution with `logits` representing the posterior over the final discrete state. Its `batch_shape` must match the HMM’s batch shape and its event shape should be scalar (a single categorical over `state_dim` categories).
- For `GaussianHMM`, the method should return a `MultivariateNormal` distribution representing the posterior over the continuous hidden state. Its `batch_shape` must equal the HMM’s batch shape, and its event shape must be `(hidden_dim,)`.
- The input `value` is the observed sequence; it must be broadcastable to the HMM’s batch plus time dimensions. The filtering computation should correctly combine initial, transition, and observation factors, summing over hidden states (for discrete) or performing Gaussian tensor contraction (for continuous) across the time dimension.

**Additional bug:**
The `expand` methods of both `DiscreteHMM` and `GaussianHMM` assign the `validate_args` flag to the public attribute `new.validate_args` instead of the internal `new._validate_args`. This causes the expanded distribution to ignore the original validation setting, leading to potential silent correctness issues. The flag should be stored in `new._validate_args` to respect the distribution’s validation contract.

**Impact:**
Users cannot cleanly chain filtering with a forecasting model, making time‑series prediction cumbersome and error‑prone. The `validate_args` bug may produce invalid shapes or log‑probabilities without raising errors when validation is expected.

Implement *HMM.filter() methods for forecasting (#2010)


We use `Binomial` and `BetaBinomial` distributions in models with discrete latent variables, and we rely on `TraceEnum_ELBO` to automatically enumerate over their domains. However, these distributions raise a `ValueError` whenever `log_prob` is evaluated on a value outside their natural support (e.g., a negative integer or an integer larger than `total_count`). During enumeration the algorithm often tries values that are outside the valid integer interval, causing the entire inference to fail.

**Expected behavior**
`log_prob` should return `-inf` for any integer input that does not satisfy the distribution’s original constraints, rather than raising an exception. The support of the distribution should be declared as `constraints.integer` so that enumeration can iterate over the entire integer domain safely. This matches the behaviour of other discrete distributions like `Categorical`, which already handle out‑of‑support values by returning `-inf`.

**Impact**
Without this change, `Binomial` and `BetaBinomial` cannot be used in models that perform automatic discrete enumeration (e.g., via `TraceEnum_ELBO`). This severely limits the applicability of these distributions in Pyro’s enumeration-based inference and makes it impractical to build models that combine these likelihoods with discrete latent variables.

**Additional requirement**
A new integer constraint seems necessary, because Pyro currently does not define an `integer` constraint that can be attached as the support of such extended distributions. The constraint should simply check that a value is integral (`value % 1 == 0`).

Add extended Binomial and BetaBinomial distributions (#2414)


Replace the graph-based downstream cost computation in `TraceGraph_ELBO` with provenance tracking. Annotate values of non‑reparameterizable sample sites with their site name using `pyro.ops.provenance.track_provenance` during tracing. Then, in the ELBO surrogate loss, collect model and guide `log_prob` terms into per‑site downstream cost tensors by consulting the provenance sets of each `log_prob_sum` (model terms added, guide entropy subtracted). This removes the need for the existing `_compute_downstream_costs` helper and should reduce gradient estimator variance by capturing exact data‑flow dependencies. Ensure that the final loss values have their provenance detached to avoid contaminating later computations.

Use provenance tracking to compute downstream costs in `TraceGraph_ELBO` (#3081)


Bug: premature key drop timer for initial key phase

After the handshake, QUIC keys start at key phase 0. When the local endpoint initiates a key update (e.g., from phase 0 to phase 1) and then receives the first packet protected with the new key phase, the code schedules the old keys to be dropped after a timeout. The problem is that the same scheduling logic fires **when no key update has occurred yet** — i.e., when the local endpoint is still on key phase 0 and receives its first packet with phase 0. Since key phase 0 represents the initial keys used immediately after the handshake, there is no older key generation to retire. Setting a drop timer in this situation risks premature removal of the valid, current keys, which would cause decryption failures for subsequent legitimate packets from the peer that use phase 0.

Acceptance criteria:
- The initial key phase (phase 0) must **never** be scheduled for timed drop. Only keys that are truly "previous" (i.e., a key phase less than the current one after a confirmed key update) should be dropped on a timer.
- When the local endpoint successfully advances to a key phase > 0 and the peer confirms by sending a packet with that phase, the old keys should still be subject to the usual drop timer.
- The overall key‑update flow (initiating update, receiving first packet with new phase, then discarding old keys) must continue to work correctly for any update from phase N to N+1 where N > 0.

The relevant state machine is in the `Open` method of the updatable AEAD, where `firstRcvdWithCurrentKey` and the `startKeyDropTimer` call govern when the old keys expire. The fix should ensure that the drop timer is only started when a genuine roll of keys has been confirmed (i.e., current key phase is greater than 0).

don't set a timer to drop keys for key phase 0


**Title:** Stream flow control incorrectly rejects duplicate final offsets

**Description**
In a QUIC connection, when a stream receives a frame with a FIN bit and an offset equal to the stream’s already known final size (i.e., a retransmission of the final offset), the library currently returns a `FinalSizeError` transport error. RFC 9000 requires that receiving the same final offset again is acceptable and must not cause an error.

Similarly, receiving a non‑final offset that is already the highest received should be silently ignored instead of causing any side effects or state changes in the flow controller.

**Acceptance criteria**
- Calling `UpdateHighestReceived(offset, final: true)` with an offset that was already accepted as final must succeed (no error).
- `UpdateHighestReceived` with a non‑final offset equal to the current highest received must succeed and be a no‑op (no window state changes, no error).
- The existing behavior of detecting genuinely inconsistent final offsets (a different final offset after one was already set) and flow‑control violations (offsets beyond the receive window) must not be altered.

**Additional note**
While fixing this, verify that other flow‑control state transitions (blocked status, send‑window size visibility) remain consistent. The `EnsureMinimumWindowSize` routine on the connection‑level controller should only request a window increase and reset its auto‑tuning epoch when the requested increment actually exceeds the current window size.

flowcontrol: rewrite tests, use require instead of Ginkgo (#4734)


The QUIC transport should attempt to raise the UDP receive buffer to a recommended size (defined by the protocol) for optimal throughput. Currently, if the OS rejects the ordinary `Setsockopt` (for example due to system-wide maximum buffer limits), connection setup returns an error and fails. We want to make this step resilient: allow the buffer increase to fail silently, and always verify the actual buffer size afterward.

On Linux, if after a normal increase the buffer is still below the desired size, a second attempt should be made using `SO_RCVBUFFORCE` (only available with `CAP_NET_ADMIN`). This forces the buffer beyond the usual system limit. No such fallback should be attempted on non-Linux systems.

In all cases, the QUIC connection must not fail solely because the desired buffer size could not be reached. The implementation should log the outcome but continue with whatever buffer size the kernel actually grants.

use SO_RCVBUFFORCE to force receive buffer increase on Linux (#3804)


supervision_value: 3
difficulty: 3
reasoning: The change threads a new public interface through the HTTP/3 response body, requiring the agent to inspect how response bodies are constructed, compare request vs response paths, locate where the underlying `quic.Connection` is available in the client, and adjust the body type and its construction without breaking existing read/close semantics. The work involves examining `body.go`, `client.go`, and the QUIC connection interfaces, then making several coordinated edits across files. This is a solid mid-range task: not trivial, not overwhelmingly broad.

---

**Feature request**

Add an HTTP/3 response-body hijacking capability that allows callers to open new QUIC streams on the same connection after receiving a response. This is needed for protocols like WebTransport that bootstrap multiplexed streams from a single HTTP/3 session.

Acceptance criteria:

1. Define a public `StreamCreator` interface in the `http3` package that mirrors the stream-creation subset of `quic.Connection`:
   - `OpenStream() (quic.Stream, error)`
   - `OpenStreamSync(context.Context) (quic.Stream, error)`
   - `OpenUniStream() (quic.SendStream, error)`
   - `OpenUniStreamSync(context.Context) (quic.SendStream, error)`
   Ensure `quic.Connection` satisfies this interface (compile-time check).

2. Define a public `Hijacker` interface in the `http3` package:
   - `StreamCreator() StreamCreator`
   - Document that it is meant to be type-asserted from an `http.Response.Body` obtained via an HTTP/3 client. Mention the WebTransport use case.

3. When a client makes an HTTP/3 request, make the resulting `http.Response.Body` implement `Hijacker` so that callers can type-assert:
   ```go
   if hj, ok := resp.Body.(http3.Hijacker); ok {
       sc := hj.StreamCreator()
       stream, _ := sc.OpenStreamSync(ctx)
       ...
   }
   ```

4. The returned `StreamCreator` must allow opening new streams on the same QUIC connection that carried the response. It must be usable independently of reading the response body; closing the body must not invalidate the stream creator.

5. The response body must continue to function as a standard `io.ReadCloser`. Request bodies must *not* implement `Hijacker`; the hijack ability is only for response bodies.

6. The change must not alter the existing public surface or behavior of the HTTP/3 client beyond the new interface availability. Existing code that does not type-assert for `Hijacker` should see no difference.

add a http3.Hijacker that allows stream creation on a QUIC session from a http.Response.Body


# Bug: Reset stream completion blocked after RESET_STREAM frame loss

When a send stream has been cancelled (cancel write), the initial RESET_STREAM frame can be lost and retransmitted. In the current code, if the original RESET_STREAM frame is lost, the retransmission works, but when the retransmitted frame is later acknowledged the stream never becomes “completed”. This means resources are not released and outstanding stream counters never reach zero, blocking normal stream lifecycle.

The root cause appears to be that `sendStreamResetStreamHandler.OnLost` does not decrement the outstanding frame counter, mirroring what `sendStreamAckHandler.OnLost` already does for stream data frames. Because the counter is not decreased upon loss, the stream remains in an open state even after the retransmitted RESET_STREAM is acked.

## Expected behaviour
When a RESET_STREAM frame is lost and a retransmitted RESET_STREAM is successfully acknowledged, the send stream should transition to completed (i.e. `numOutstandingFrames` should reach zero, triggering a call to `sender.onStreamCompleted`).

## Steps to reproduce (conceptual)
1. Create a send stream.
2. Cancel the write side (`CancelWrite`).
3. Arrange for the initial RESET_STREAM frame to be lost.
4. Observe that a retransmitted RESET_STREAM frame is sent.
5. Acknowledge the retransmitted frame.
6. The stream still has outstanding frames and does not complete.

## Fix guidance
Locate the loss handler for RESET_STREAM frames inside `send_stream.go` and ensure it properly decrements the outstanding frame counter. No other changes should be needed; the corresponding ack handler already handles stream completion when the counter reaches zero.

## Acceptance criteria
- After a lost-then-retransmitted RESET_STREAM frame is acked, the stream completes.
- No regression in existing cancellation / reset retransmission scenarios.

fix accounting for lost RESET_STREAM frames in the send stream (#4804)


**Title:** qlog does not record QUIC key update events

**Description:**

The qlog trace records `key_updated` events for TLS key derivation (trigger `"tls"`) but never emits events for 1‑RTT key updates that occur at the QUIC layer. This makes it impossible to correlate key‑update timing with other transport events in the trace, which is important for debugging performance and loss‑recovery interactions.

Key updates happen in two ways:
1. **Remote update:** when a peer rotates its keys and the receiver detects the new key phase.
2. **Local update:** when the current endpoint decides to rotate its own keys (for example, after the key‑update interval).

In both cases the code already logs a debug message (e.g. `“Peer updated keys to …”`), proving the events are detected, but no corresponding qlog event is produced.

**Expected behavior:**

Whenever a 1‑RTT key rotation is detected or initiated, the qlog tracer should emit a `key_updated` event with:
- `trigger`: `"remote_update"` for a peer‑triggered rotation; `"local_update"` for a self‑initiated rotation.
- `key_type`: both `"server_1rtt_secret"` and `"client_1rtt_secret"` (the rotation affects the 1‑RTT encryption keys for both directions simultaneously).
- `generation`: the current key phase number (`protocol.KeyPhase`).

**Constraints / safety:**
- The qlog tracer may be `nil` (qlog disabled). All new logging calls must guard against `nil`.
- The distinction between remote and local triggers must be preserved, so that the trace accurately reflects who initiated the rotation.
- The existing `UpdatedKeyFromTLS` method should remain unchanged; the QUIC‑layer events are a separate concern.

**Impact:**

Without this, qlog users have no visibility into when and why key rotations happen, making it harder to diagnose issues that are sensitive to key‑update timing (e.g. packet loss around key changes, stalls, or interactions with the recovery controller).


Make the initial packet size used for outgoing QUIC packets configurable via a new `Config` field, instead of being hardcoded to `protocol.InitialPacketSize`. The new field should be a `uint16` with a zero value meaning “use the current default.” Values below 1200 must be clamped to 1200, and values above `protocol.MaxPacketBufferSize` must be clamped to that maximum. Ensure the configured size is applied consistently across client and server connection initialization, including initial MTU discovery and payload size estimation. The initial packet size must be used for the very first UDP packet sent.

make the initial packet size configurable (#4503)


supervision_value: 4
difficulty: 3
reasoning: The task requires understanding state management for QUIC stream cancellation, connecting the `CancelWrite` flow with acknowledgement and loss callbacks, and reasoning about completion invariants. It is not a trivial one-line guard; the fix involves tracking cancellation across three methods and updating internal counters. The scope is still bounded to one file and a set of related tests, making it a strong P2T candidate with medium difficulty.

---

**Title:** send stream: cancellation with outstanding STREAM frames fails to complete the stream and panics on later frame events

**Description:**

When `CancelWrite(errorCode)` is called on a QUIC send stream while unacknowledged STREAM frames are still in flight (i.e., they have been popped and sent but not yet acknowledged), two problems occur:

1. **Stream never completes** – The stream is not signalled as “completed” after cancellation, even though the application has explicitly terminated it. Downstream completion‑tracking (e.g., connection‑level stream cleanup, session shutdown logic) therefore never receives the `onStreamCompleted` notification, which can stall connection shutdown or leak stream resources.

2. **Later ACK/loss events cause panics or corruption** – If those in‑flight frames are subsequently acknowledged or detected as lost, the stream’s internal outstanding‑frame counter can underflow (panic: `numOutstandingFrames negative`). The retransmission queue may also be populated with frames for a stream that should no longer be retransmitting, leading to unexpected behaviour.

**Expected behaviour:**

- After `CancelWrite` the stream must immediately be considered finished (since it is being forcefully terminated).
- The stream’s completion callback (`onStreamCompleted`) should be invoked exactly once, as part of the cancellation.
- Any acknowledgements or loss notifications for frames that were sent before the cancellation must be ignored; they must not modify the stream’s state, trigger retransmissions, or affect outstanding‑frame accounting.
- The stream should not produce any further data frames for retransmission after being cancelled.

**Steps to reproduce (conceptually):**

1. Open a send stream and write data such that a STREAM frame is popped and placed into the ack‑handler (in‑flight).
2. Before that frame is acknowledged, call `CancelWrite` on the stream.
3. Later, simulate a loss or acknowledgement for that frame.
4. Observe that a panic (negative outstanding frames) occurs, or that `onStreamCompleted` was never called, or that the stream attempts to retransmit data it no longer owns.

**Impact:**

- Crashes (panics) in production when connections are cancelled at the right (wrong) moment.
- Streams and associated resources (flow control windows, callbacks) may never be cleaned up, leading to slow resource exhaustion or hanging shutdowns.
- The `Stream` interface contract is violated: cancellation should be final and immediate, but here it leaves the stream in a half‑alive state that reacts to network events.

complete a stream that has frames in flight if it is canceled


ackhandler: store skipped packet numbers separately in sentPacketHistory

Currently, `sentPacketHistory` tracks packet numbers that were intentionally skipped
(not sent) by inserting placeholder `packet` entries that carry a `skippedPacket: true`
flag. That representation has two problems:

1. Many callers that iterate over `Packets()` must remember to filter out skipped
   entries. This manifests as repeated `!p.skippedPacket` guards in loss detection,
   retransmission, ACK validation, path migration, and `outstanding()` – and it is
   easy for a new call site to miss the check.

2. Skipped numbers are not “packets” in any meaningful sense (they have no sent
   time, no frames, no encryption level). Storing them inside the regular packet
   list makes the data model misleading and complicates history maintenance (e.g.
   cleanup / removal of leading entries).

The ask:

Refactor `sentPacketHistory` so that skipped packet numbers are stored separately
from actual packets.  The `Packets()` iterator should yield only real packets;
skipped numbers must not appear in it.  The history’s `Len()` should reflect only
real packets as well.

Callers that currently inspect `p.skippedPacket` should instead not see skipped
entries at all, and the special-casing should be removed from:

- `packet.outstanding()`
- `detectAndRemoveAckedPackets` (the per-packet check for skipped → protocol violation)
- `detectLostPackets`
- `DropPackets`, `ResetForRetry`, `MigratedPath`

Provide a separate, stable way to obtain the list of skipped packet numbers so that
the ACK handler can still detect the protocol violation “peer acknowledged a skipped
packet number” (RFC 9000 §13.2.1) without relying on the main packet iterator.

When a real packet is removed from the history, any skipped numbers that are no
longer needed (e.g. because they precede the new first real packet) should be
garbage-collected, keeping the internal state minimal.  The implementation must
continue to enforce sequential packet-number use, and must preserve the existing
behaviour of the history’s “first packet number” tracking and `cleanupStart`.

The goal is a correct, minimal refactoring that removes every `skippedPacket`-based
check from the packet type and from the callers, while giving each site exactly the
data it needs.

ackhandler: store skipped packet numbers separately (#5314)


**Title:** Add tracer hook for sent Version Negotiation packets

**Problem:** The `logging.Tracer` interface does not provide a way to observe outgoing QUIC Version Negotiation packets. A Version Negotiation packet is sent by a server when it receives a client connection with an unsupported version. The packet carries a list of supported version numbers and uses connection IDs that can be of arbitrary length (up to 255 bytes). Currently, the only packet-sent hook is `SentPacket(net.Addr, *Header, ByteCount, []Frame)`, which is designed for packets that have a regular QUIC header and that are sent within an established connection. Version Negotiation packets are sent before any connection is established and do not fit that signature, making it impossible for a `Tracer` implementation to log or monitor all outgoing QUIC packets.

**Expected behavior:** The `Tracer` interface should expose a new method dedicated to Version Negotiation packets so that tracing tools can capture every packet type the stack emits. The method signature should include:
- the remote network address,
- the destination and source connection IDs (which may be `ArbitraryLenConnectionID`, a type currently internal to the protocol layer),
- the list of advertised QUIC version numbers.

All concrete `Tracer` implementations shipped in the library must satisfy the extended interface:
- The null tracer should provide a no-op implementation.
- The multiplexing tracer should fan the call out to all underlying tracers.
- The qlog tracer (and any other bundled tracer) must compile correctly after the interface change; it may adopt the null tracer’s behavior for now if proper qlog wiring is deferred.

**Additional requirement – type visibility:** The `logging` package currently re-exports `ConnectionID` but not `ArbitraryLenConnectionID`. Because the new tracer method needs to receive arbitrary-length connection IDs, the `logging` package must also export `ArbitraryLenConnectionID` as a type alias so that external tracer implementations can reference it.

**Why it matters:** Without this hook, any user-provided tracer that aims to give complete observability of the QUIC transport layer will have a blind spot for the very first packet the server may send. This is especially problematic for debugging version negotiation issues, network analysis tools, and logging infrastructures that expect a complete packet trace.

add a tracer function to trace sending of Version Negotiation packets


The `qlogwriter/jsontext` encoder is used to generate qlog output and deliberately keeps its feature set minimal.  It already supports strings, numbers, booleans, objects, and arrays, but there is currently no way to output a JSON `null` literal.

We need the ability to emit `null`.  This is important because qlog schemas include optional fields that are explicitly set to `null`, and our encoder must be able to represent them.

**Requested change**

* Add a new `Null` token (similar to `True` and `False`), so callers can write `enc.WriteToken(jsontext.Null)`.
* When the encoder processes this token, it must output the bytes `null`, correctly handle comma insertion when inside objects or arrays, and update the internal after-value state exactly as it does for other value tokens (`Bool`, `Int`, `Float`, `String`, etc.).

No new configuration or grammar rules are needed.  The existing `encodeValue` helper in the test file should also be extended to map Go `nil` to the new `Null` token, ensuring fuzz and comprehensive tests continue to work correctly.

jsontext: add support for encoding null (#5371)


Add a new exported field `EnableDatagrams bool` to `http3.Server`.
When true, two things must happen:

1. The underlying QUIC connection must be created with datagrams enabled: the server must set `quic.Config.EnableDatagrams = true` on the configuration it passes to the QUIC listener.
   **Important:** do not mutate the caller’s `Server.QuicConfig` in place. If a non-nil `*quic.Config` is provided, clone it before setting the flag; if nil, allocate a fresh zero config.

2. The HTTP/3 SETTINGS frame that the server writes on the control stream must include the `SETTINGS_H3_DATAGRAM` parameter (value 0x33) to advertise datagram support, following draft-ietf-masque-h3-datagram (e.g. the 02 draft).

Everything else (regular HTTP/3 serving) should remain unchanged.

add a http3.Server.EnableDatagram option


```
# http3: implement trailer header validation

## Problem
When trailers are received on an HTTP/3 stream, `parseTrailers` currently
does not validate the individual header fields. A TODO in the code
indicates that validation is missing.

## RFC requirements
Per RFC 9114, trailer fields must be validated against the same rules as
regular header fields, with the additional restriction that the trailer
field name must be acceptable as a trailer according to
`httpguts.ValidTrailerHeader` (which excludes headers like
`Content-Length`, `If-Match`, etc.).

More concretely, the validation must reject:
- pseudo headers in trailers
- field names that are not lowercase
- field names or values that fail `httpguts.ValidHeaderFieldName` / `httpguts.ValidHeaderFieldValue`
- connection-specific headers (e.g., `Connection`, `Keep-Alive`, …)
- `TE` headers with a value other than `"trailers"`
- any field name that is *not* a valid trailer per `httpguts.ValidTrailerHeader`

Additionally, the uncompressed trailer field list must be checked against
the header size limit (`maxHeaderBytes`) using the same accounting as for
regular headers (name length + value length + 32 bytes overhead per field,
see RFC 9114 §4.2.2).

## Desired behavior
- Trailer parsing must return an error for any of the violations above.
  The error should *not* be a QPACK error; it must be a distinct error
  that the transport layer can map to a stream-level reset or
  `H3_FRAME_UNEXPECTED` / `H3_MESSAGE_ERROR` as appropriate.
- The call path that triggers trailer decoding (`decodeTrailers`) must
  pass the connection-wide `maxHeaderBytes` limit down into the trailer
  parser so that oversized trailers are rejected.
- All existing regular-header validation logic should be reused rather
  than duplicated; any refactoring is fine as long as the behavior is
  identical for headers and becomes strictly stricter (correct) for
  trailers.
- A trailer with a single valid field (e.g., `foo: 42`) must still be
  accepted and parsed into the `http.Header` map as before.
```

http3: implement trailer validation logic (#5642)


When a QUIC server is amplification-limited (hasn’t received enough bytes from the client to send freely), we deliberately suppress the loss detection timer.
When the server later becomes unblocked after receiving additional bytes, the loss detection timer is not being restarted. This causes loss detection to remain disabled indefinitely, which can stall recovery.

Fix this by restarting the loss detection timer whenever the server transitions from being amplification-limited to no longer being limited.

restart the loss detection timer when the server becomes unblocked


Add the QUIC version 2 constant (`0x709a50c4`). Include it in the list of supported versions and make sure the version’s string representation returns `"v2"`.


```
Title: 0-RTT resumption incorrectly rejected when server increases its max stream limit

Observed behavior:
When a QUIC server that previously issued a session ticket is restarted with a higher MaxIncomingStreams (or MaxIncomingUniStreams) value, clients that attempt to resume the session using 0-RTT are refused. The handshake falls back to a full 1-RTT exchange even though the server’s new limit is only more permissive than what was recorded in the ticket.

Expected behavior:
Increasing the server’s maximum bidirectional or unidirectional stream limit should not invalidate 0-RTT resumption. The client is still bound by the lower limit stored in its session ticket, and that limit remains compliant with the server’s new (higher) cap. Conversely, a decrease in the stream limit must continue to reject 0-RTT, because the client might believe it can open more streams than the server now permits.

The relevant transport parameters stored inside the session ticket are:
- initial_max_streams_bidi (server → client, corresponding to MaxIncomingStreams)
- initial_max_streams_uni (server → client, corresponding to MaxIncomingUniStreams)

Other transport parameters (initial_max_data, initial_max_stream_data_*, active_connection_id_limit) should still be required to match exactly, since a change there would affect credit or connection management in ways that are not safe for 0-RTT.

Impact:
This makes it impossible to scale up stream limits across server restarts without breaking 0-RTT for clients that hold older tickets, harming connection latency.

Request:
Update the 0-RTT eligibility check so that stream-limit transport parameters are accepted if they have *increased or stayed the same* compared to the values stored in the session ticket. All other transport parameters should continue to be accepted only when they are exactly equal.
```

allow 0-RTT resumption if the server's stream limit was increased