Skip to content

Fix reverse-mode gradients for Dirichlet batches - #829

Open
AHMETHAKANBEZIR1 wants to merge 1 commit into
HIPS:masterfrom
AHMETHAKANBEZIR1:fix/dirichlet-batch-gradients
Open

AHMETHAKANBEZIR1 wants to merge 1 commit into
HIPS:masterfrom
AHMETHAKANBEZIR1:fix/dirichlet-batch-gradients

Conversation

@AHMETHAKANBEZIR1

Copy link
Copy Markdown

SciPy evaluates Dirichlet densities for a matrix whose columns are samples, but the reverse-mode rules broadcast the concentration vector against the sample axis and do not reduce concentration gradients across samples. For example, differentiating a summed logpdf with three components and two samples raises a broadcasting error even though density evaluation succeeds.

Compute the local derivatives in sample-major order, restore the original coordinate layout, and use the existing unbroadcast_f helper to accumulate gradients for shared concentrations. This covers batches with all simplex components supplied; support for SciPy's abbreviated last-component input is unchanged.

Validation:

  • Six new weighted batch cases (one, two, and three samples; PDF and log-PDF), checking both arguments through second-order reverse-mode differentiation: all six fail on unchanged master and pass with the patch. Four existing single-sample controls pass throughout.
  • Full test suite: python -m pytest -n 2 — 651 passed, 14 warnings on Python 3.12.14, NumPy 2.5.4, SciPy 1.18.1.
  • Configured pre-commit hooks and git diff --check pass. Forward-mode support is unchanged.

AI disclosure: OpenAI Codex autonomously prepared the code, tests, and this description. Validation and an additional AI review were performed; there has been no independent human review.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant