10000 Add documentation for recursive type demo by sunxd3 · Pull Request #588 · chalk-lab/Mooncake.jl · GitHub
[go: up one dir, main page]
More Web Proxy on the site http://driver.im/
Skip to content

Add documentation for recursive type demo #588

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 28 commits into from
May 29, 2025
Merged

Add documentation for recursive type demo #588

merged 28 commits into from
May 29, 2025

Conversation

sunxd3
Copy link
Collaborator
@sunxd3 sunxd3 commented May 27, 2025 •
edited by yebai
Loading

Partial solution to #428

@sunxd3 sunxd3 marked this pull request as ready for review May 27, 2025 12:35
@sunxd3 sunxd3 marked this pull request as draft May 27, 2025 12:35
@codecov Codecov
Copy link
codecov bot commented May 27, 2025

Codecov Report

All modified and coverable lines are covered by tests ✅

📢 Thoughts on this report? Let us know!

@github-actions GitHub Actions
Copy link
Contributor

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

@github-actions GitHub Actions
Copy link
Contributor
github-actions bot commented May 27, 2025 •
edited
Loading

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 │  Zygote │ ReverseDiff │  Enzyme │
│                     String │   String │   String │  String │      String │  String │
├────────────────────────────┼──────────┼──────────┼─────────┼─────────────┼─────────┤
│                   sum_1000 │ 100.0 ns │      1.8 │     1.0 │        5.61 │    8.21 │
│                  _sum_1000 │ 941.0 ns │     6.59 │  1550.0 │        33.5 │    1.09 │
│               sum_sin_1000 │  6.56 μs │     2.22 │    1.67 │        11.0 │    2.22 │
│              _sum_sin_1000 │  5.37 μs │     2.59 │   263.0 │        13.0 │    2.43 │
│                   kron_sum │ 269.0 μs │     48.6 │     5.8 │       239.0 │    13.8 │
│              kron_view_sum │ 322.0 μs │     42.6 │    11.4 │       227.0 │    9.59 │
│      naive_map_sin_cos_exp │  2.17 μs │     2.13 │ missing │        7.21 │    2.29 │
│            map_sin_cos_exp │   2.1 μs │     2.42 │    1.59 │         6.2 │    2.89 │
│      broadcast_sin_cos_exp │  2.29 μs │     2.21 │    2.31 │        1.44 │    2.21 │
│                 simple_mlp │ 209.0 μs │     5.82 │    1.67 │        10.2 │    3.23 │
│                     gp_lml │ 242.0 μs │     8.97 │    3.86 │     missing │    4.83 │
│ turing_broadcast_benchmark │  1.98 ms │     3.54 │ missing │        25.7 │ missing │
│         large_single_block │ 381.0 ns │     4.49 │  4190.0 │        31.3 │    2.23 │
└────────────────────────────┴──────────┴──────────┴─────────┴─────────────┴─────────┘

@yebai
Copy link
Member
yebai commented May 27, 2025 •
edited
Loading

The "known limitations" doc page (circular reference section) also contains some related but very brief explanations.

penelopeysm
penelopeysm reviewed May 27, 2025
View reviewed changes
Copy link
Collaborator
@penelopeysm penelopeysm left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Very nice overview --- I'm just listing some things that would be icing on the cake, really.

docs/src/custom_tangent_type.md Outdated
Comment on lines 61 to 64

function TangentForA{Tx}(nt::@NamedTuple{x::Tx, a::Union{Mooncake.NoTangent, TangentForA{Tx}}}) where Tx
return new{Tx}(nt.x, nt.a)
end
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this method necessary, or just for convenience?

If it's just for convenience, I think I'd prefer to leave it out, otherwise a reader may assume that it's part of the interface that's needed. If it's needed, then that should be mentioned IMO.

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not necessary for this PR, but necessary to pass TestUtils.test_data. Removing.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Ah ok in that case actually it could be moved below too!

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sorry, said too soon. It's indeed necessary here.

docs/src/custom_tangent_type.md Outdated

### 4.2. `fdata_type` and `rdata_type`

Define the types of forward and reverse data for `TangentForA`. In this example, since both `A` and `TangentForA` are mutable, all updates can be done in-place, so the `fdata` is the tangent itself and `rdata` is [`NoRData`](@ref). In other cases, you may need to split these more carefully.
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Do you need to declare methods for fdata and rdata?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a sentence for that.

docs/src/custom_tangent_type.md Outdated
f1(a::A) = 2.0 * a.x
```

When you try to differentiate this, Mooncake will complain it lacks an `rrule!!` for `lgetfield`. Implement it:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Maybe a brief mention of what lgetfield is might be useful

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added a ref link

docs/src/custom_tangent_type.md Outdated
Comment on lines 189 to 190
## 6. From "It Works!" to Passing [`test_data`](@ref)

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would it be possible to have a snippet of using test_data? (It's ok if it fails -- that's the point of this section, but I think it'd be nice to have a brief example of 'this is broken, below is how you fix it'.)

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I added all the code need to pass the test. Probably should live somewhere else, but okay for now.

@sunxd3 sunxd3 mentioned this pull request May 27, 2025
Open
@sunxd3 sunxd3 marked this pull request as ready for review May 27, 2025 16:55
AstitvaAggarwal
AstitvaAggarwal reviewed May 27, 2025
View reviewed changes
@AstitvaAggarwal
Copy link
Collaborator

This work on the documentation is really good, just added some comment for minor changes. beyond that it looks comprehensive and GTM.

sunxd3 and others added 4 commits May 29, 2025 13:22
@sunxd3 @AstitvaAggarwal
Update docs/src/custom_tangent_type.md
ebd0dcf 
Co-authored-by: Astitva Aggarwal <84859349+AstitvaAggarwal@users.noreply.github.com>
Signed-off-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
@sunxd3 @AstitvaAggarwal
Update docs/src/custom_tangent_type.md
4e7cab0 
Co-authored-by: Astitva Aggarwal <84859349+AstitvaAggarwal@users.noreply.github.com>
Signed-off-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
@sunxd3 @AstitvaAggarwal
Update docs/src/custom_tangent_type.md
6e95c28 
Co-authored-by: Astitva Aggarwal <84859349+AstitvaAggarwal@users.noreply.github.com>
Signed-off-by: Xianda Sun <5433119+sunxd3@users.noreply.github.com>
@sunxd3
address comments
1f0117b
@sunxd3
Copy link
Collaborator Author
sunxd3 commented May 29, 2025

Thanks both for the review! Really helpful.

@yebai
Copy link
Member
yebai commented May 29, 2025

@sunxd3, you need to (re-)format the code to address the "quality" CI failures.

yebai
yebai reviewed May 29, 2025
View reviewed changes
@yebai
Copy link
Member
yebai commented May 29, 2025

I moved the new doc page to the "Developer" section. I'm happy to merge!

yebai
yebai reviewed May 29, 2025
View reviewed changes
@yebai
Update docs/src/known_limitations.md
dc1280c 
Signed-off-by: Hong Ge <3279477+yebai@users.noreply.github.com>
@sunxd3
Copy link
Collaborator Author
sunxd3 commented May 29, 2025

thanks @yebai

@yebai
Update known_limitations.md
cff8e24 
Signed-off-by: Hong Ge <3279477+yebai@users.noreply.github.com>
@yebai yebai merged commit 5f1e8aa into main May 29, 2025
76 of 77 checks passed
@yebai yebai deleted the recursive_type_demo branch May 29, 2025 16:18
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@yebai yebai yebai left review comments

@AstitvaAggarwal AstitvaAggarwal AstitvaAggarwal left review comments

@penelopeysm penelopeysm penelopeysm left review comments

Assignees

@sunxd3 sunxd3

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
4 participants
@sunxd3 @yebai @AstitvaAggarwal @penelopeysm
0