feat: add gradient fallback helpers (tesseract_core.runtime.experimental.vjp_from_jacobian, ...) for deriving AD endpoints from each other

[andrinr]

Relevant issue or PR

Resolves #465 when merged

Description of changes

Adds four experimental helpers in tesseract_core.runtime.experimental for deriving
missing AD endpoints from ones already implemented:

• jvp_from_jacobian / vjp_from_jacobian — derive JVP or VJP from an existing jacobian endpoint by contracting the Jacobian matrix with the tangent / cotangent vector
• jacobian_from_jvp — materialise the full Jacobian by probing with N one-hot tangents (N = input elements); output shapes are inferred from the first probe, so no extra apply call is needed
• jacobian_from_vjp — materialise the full Jacobian by probing with M one-hot cotangents (M = output elements); takes an eval_fn argument (either apply or the cheaper abstract_eval) to determine output shapes upfront

All helpers live in tesseract_core/runtime/ad_endpoint_derivation.py and are re-exported through tesseract_core.runtime.experimental.

Testing done

• tests/runtime_tests/test_autodiff_fallbacks.py — unit tests parametrised over two nonlinear functions (wide R⁴→R³ and tall R²→R⁴); jacobian_from_vjp additionally parametrised over apply vs abstract_eval; round-trip tests verify VJP→Jacobian→JVP and JVP→Jacobian→VJP chains
• examples/univariate_adfallbacks — e2e Rosenbrock variant using jvp_from_jacobian / vjp_from_jacobian; covers test_apply, test_jacobian_wrt_x, test_jvp, test_vjp, and check_gradients; registered in tests/endtoend_tests/test_examples.py with check_gradients=True

[codecov]

Codecov Report

❌ Patch coverage is 82.45614% with 10 lines in your changes missing coverage. Please review.
✅ Project coverage is 77.06%. Comparing base (9ae8e62) to head (dda9a23).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
tesseract_core/runtime/ad_endpoint_derivation.py	82.14%	4 Missing and 6 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #511      +/-   ##
==========================================
+ Coverage   67.74%   77.06%   +9.32%     
==========================================
  Files          31       32       +1     
  Lines        4291     4348      +57     
  Branches      705      718      +13     
==========================================
+ Hits         2907     3351     +444     
+ Misses       1146      710     -436     
- Partials      238      287      +49     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[dionhaefner]

I'm wondering if we could address this in a less magical way by providing functions like tesseract_core.runtime.experimental.jvp_from_jacobian (similar to the new FD functionality), so users can opt-out / in explicitly?

[andrinr]

@dionhaefner good idea! I think the only advantage of this approach is that we can add such endpoints at serve time and do not the rebuild the tesseract. But I also see arguments against this, in a way it changes the predefined contract of a tesseract which is a bit of an antipattern.

I am leaning towards the approach you just proposed.

[dionhaefner]

Yes I think the only time this is worse is when there's someone wanting to use a pre-built Tesseract image in a workflow that requires VJPs but the Tesseract only has Jacobian defined. OTOH we can still enable things globally if the experimental feature proves really useful.

[PasteurBot]

Benchmark Results

Benchmarks use a no-op Tesseract to measure pure framework overhead.

🚀 0 faster, ⚠️ 0 slower, ✅ 36 unchanged

✅ No significant performance changes detected.

Full results

Benchmark	Baseline	Current	Change	Status
api/apply_100	0.794ms	0.809ms	+1.9%	✅
api/apply_10,000	0.838ms	0.831ms	-0.8%	✅
api/apply_1,000,000	2.342ms	2.281ms	-2.6%	✅
cli/apply_100	1750.616ms	1798.731ms	+2.7%	✅
cli/apply_10,000	1745.843ms	1811.233ms	+3.7%	✅
cli/apply_1,000,000	1786.518ms	1784.299ms	-0.1%	✅
decoding/base64_100	0.030ms	0.031ms	+3.5%	✅
decoding/base64_10,000	0.119ms	0.120ms	+0.9%	✅
decoding/base64_1,000,000	9.109ms	9.110ms	+0.0%	✅
decoding/binref_100	0.194ms	0.196ms	+1.1%	✅
decoding/binref_10,000	0.214ms	0.212ms	-1.2%	✅
decoding/binref_1,000,000	1.192ms	1.168ms	-2.0%	✅
decoding/json_100	0.034ms	0.034ms	-0.5%	✅
decoding/json_10,000	0.949ms	0.932ms	-1.8%	✅
decoding/json_1,000,000	109.225ms	107.356ms	-1.7%	✅
encoding/base64_100	0.041ms	0.040ms	-2.8%	✅
encoding/base64_10,000	0.352ms	0.351ms	-0.1%	✅
encoding/base64_1,000,000	37.804ms	38.783ms	+2.6%	✅
encoding/binref_100	0.288ms	0.287ms	-0.4%	✅
encoding/binref_10,000	0.341ms	0.371ms	+8.8%	✅
encoding/binref_1,000,000	5.737ms	5.753ms	+0.3%	✅
encoding/json_100	0.091ms	0.090ms	-1.4%	✅
encoding/json_10,000	5.313ms	5.242ms	-1.3%	✅
encoding/json_1,000,000	546.072ms	544.166ms	-0.3%	✅
http/apply_100	4.488ms	4.530ms	+0.9%	✅
http/apply_10,000	5.938ms	5.959ms	+0.4%	✅
http/apply_1,000,000	161.385ms	161.577ms	+0.1%	✅
roundtrip/base64_100	0.080ms	0.082ms	+2.2%	✅
roundtrip/base64_10,000	0.490ms	0.485ms	-0.8%	✅
roundtrip/base64_1,000,000	48.277ms	48.137ms	-0.3%	✅
roundtrip/binref_100	0.506ms	0.509ms	+0.6%	✅
roundtrip/binref_10,000	0.569ms	0.573ms	+0.6%	✅
roundtrip/binref_1,000,000	6.762ms	6.815ms	+0.8%	✅
roundtrip/json_100	0.130ms	0.136ms	+4.5%	✅
roundtrip/json_10,000	6.073ms	6.241ms	+2.8%	✅
roundtrip/json_1,000,000	626.493ms	655.613ms	+4.6%	✅

• Runner: Linux 6.14.0-1017-azure x86_64

[dionhaefner]

@andrinr Please adjust the PR description and title (so we get a meaningful entry in the changelog).

[dionhaefner]

Looks great, thanks!