Try to use count_allocs instead of @allocations

[Technici4n]

Maybe this works? Unfortunately the issue cannot be reproduced locally so I have to rely on github to test it...

[github-actions]

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

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 │ 100.0 ns │      1.9 │         1.9 │     1.1 │        5.61 │   8.31 │
│                  _sum_1000 │ 941.0 ns │     6.75 │        1.01 │  1450.0 │        34.4 │   1.07 │
│               sum_sin_1000 │  6.56 μs │     2.52 │        1.41 │    1.65 │        10.5 │   2.19 │
│              _sum_sin_1000 │   5.3 μs │     3.06 │        2.18 │   268.0 │        13.2 │   2.47 │
│                   kron_sum │ 219.0 μs │     62.8 │        3.69 │    20.7 │       301.0 │   11.7 │
│              kron_view_sum │ 331.0 μs │     40.7 │        3.26 │    10.8 │       214.0 │   10.0 │
│      naive_map_sin_cos_exp │  2.17 μs │     2.34 │        1.41 │ missing │        7.39 │   2.36 │
│            map_sin_cos_exp │   2.1 μs │     2.75 │        1.46 │    1.61 │         7.8 │   2.54 │
│      broadcast_sin_cos_exp │  2.28 μs │     2.38 │        1.37 │    2.41 │        1.46 │    2.2 │
│                 simple_mlp │ 161.0 μs │     6.84 │        2.97 │    1.83 │        12.2 │   3.33 │
│                     gp_lml │ 185.0 μs │     10.6 │        2.45 │    4.44 │     missing │   5.27 │
│ turing_broadcast_benchmark │  1.75 ms │     4.85 │        3.51 │ missing │        26.5 │    2.7 │
│         large_single_block │ 380.0 ns │     4.53 │        2.03 │  4360.0 │        31.0 │   2.27 │
└────────────────────────────┴──────────┴──────────┴─────────────┴─────────┴─────────────┴────────┘

[yebai]

The rules for sum, rules for randn and randexp appear to have a genuine allocation on Julia 1.12. However, it is not exactly whether it is from the Base.sum or Mooncake rules.

See, also, #714 (comment)

[yebai]

I don’t think this is needed anymore — I’ve added the changes directly in #714.