Skip to content

Add rrule for BridgeStan's log_density function#1006

Draft
nsiccha wants to merge 3 commits intomainfrom
ns/bridgestan
Draft

Add rrule for BridgeStan's log_density function#1006
nsiccha wants to merge 3 commits intomainfrom
ns/bridgestan

Conversation

@nsiccha
Copy link
Copy Markdown
Collaborator

@nsiccha nsiccha commented Feb 25, 2026

Adds a MooncakeBridgeStanExt package extension providing rrule!! for BridgeStan.log_density.

BridgeStan's log_density calls into compiled Stan C++ code, so Mooncake cannot differentiate through it generically. This extension uses the existing log_density_gradient! API — which Stan already computes analytically — for an exact, efficient reverse-mode rule.

  • tangent_type(StanModel) = NoTangent: StanModel holds C pointers Mooncake must not recurse into
  • Two rrule!! methods: positional log_density(sm, q) and kwarg log_density(sm, q; propto, jacobian)
  • Tests cover both call forms across all propto/jacobian combinations

@github-actions
Copy link
Copy Markdown
Contributor

Mooncake.jl documentation for PR #1006 is available at:
https://chalk-lab.github.io/Mooncake.jl/previews/PR1006/

@codecov
Copy link
Copy Markdown

codecov bot commented Feb 25, 2026

Codecov Report

❌ Patch coverage is 0% with 20 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
ext/MooncakeBridgeStanExt.jl 0.00% 20 Missing ⚠️

📢 Thoughts on this report? Let us know!

@github-actions
Copy link
Copy Markdown
Contributor

Performance Ratio:
Ratio of time to compute gradient and time to compute function.
Warning: results are very approximate! See here for more context.

┌────────────────────────────┬──────────┬──────────┬─────────────┬─────────┬─────────────┬────────┐
│                      Label │   Primal │ Mooncake │ MooncakeFwd │  Zygote │ ReverseDiff │ Enzyme │
│                     String │   String │   String │      String │  String │      String │ String │
├────────────────────────────┼──────────┼──────────┼─────────────┼─────────┼─────────────┼────────┤
│                   sum_1000 │ 110.0 ns │     1.91 │        1.92 │    1.09 │        5.46 │   10.1 │
│                  _sum_1000 │ 951.0 ns │     6.75 │        1.01 │  1620.0 │        33.5 │   1.09 │
│               sum_sin_1000 │  6.57 μs │     3.19 │        1.38 │    1.77 │        10.8 │   1.89 │
│              _sum_sin_1000 │   5.1 μs │     3.21 │        2.28 │   318.0 │        13.8 │   2.57 │
│                   kron_sum │ 436.0 μs │     6.96 │        2.93 │    4.45 │       173.0 │   10.8 │
│              kron_view_sum │ 427.0 μs │     8.19 │        4.17 │    17.5 │       175.0 │   5.31 │
│      naive_map_sin_cos_exp │  2.17 μs │     2.48 │        1.44 │ missing │        7.19 │   2.32 │
│            map_sin_cos_exp │  2.14 μs │     2.87 │        1.51 │    1.75 │        6.24 │   2.93 │
│      broadcast_sin_cos_exp │   2.3 μs │     2.51 │        1.43 │    4.57 │        1.46 │   2.21 │
│                 simple_mlp │ 183.0 μs │     6.56 │        2.84 │    1.71 │        13.6 │    3.1 │
│                     gp_lml │ 201.0 μs │     10.3 │        2.46 │    4.43 │     missing │   8.04 │
│ turing_broadcast_benchmark │  1.97 ms │     4.09 │        3.06 │ missing │        30.5 │   2.05 │
│         large_single_block │ 391.0 ns │     4.43 │         2.0 │  4500.0 │        30.5 │    2.2 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

nsiccha and others added 2 commits March 16, 2026 11:53
- Apply JuliaFormatter blue style to ext/MooncakeBridgeStanExt.jl and test/ext/bridgestan/bridgestan.jl
- Add bridgestan to the ext CI test matrix so coverage is actually collected

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
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