Forward mode mutable struct support #219

oxinabox · 2023-09-28T09:02:21Z

Closes #189

Requires

test/forward_mutation.jl

oxinabox · 2023-12-29T12:28:13Z

Changes in the ChainRulesCore PR make zero_tangent more aggressive at making structural zeros.
or NoTangents for things without tangent spaces (in particular if they don't have any fields)
This is good because it should fix up a bunch of our type inference issues, since were already using zero_tangent everywhere in this PR.
A complicating factor is we also need to fix up where we use ZeroTangent for thigns that actually have no tangent space -- which isn't wrong as such but if we don't care either way we should be disaptching on AbstractZero.
Which this PR now does.

current status is WIP i know it is broken, and will fix after the long weekend.
I suspect one bit of it being broken is we need #236

oxinabox · 2024-01-01T07:05:43Z

I merged #236 in here because i was running into code that was hitting the exact case that fixes in the tests.
Which does show our code path is now significantly different, probably because of the structural zeros.

src/codegen/forward_demand.jl

oxinabox · 2024-01-12T09:37:42Z

src/codegen/forward_demand.jl

@@ -362,6 +363,6 @@ function forward_diff!(interp::ADInterpreter, ir::IRCode, src::CodeInfo, mi::Met
    rt = CC._ir_abstract_constant_propagation(interp, irsv)

    ir = compact!(ir)
-
+    


Suggested change

src/stage1/recurse_fwd.jl

oxinabox · 2024-01-22T16:53:03Z

src/stage1/recurse_fwd.jl

+            # It is a little dubious using StructuralTangent{B} for >1st order, but it is isomorphic.
+            # Just watch out for order mixing bugs.


I am not even sure it is actually dubious. Further thought has convinced me its the right thing to do.
Ironically (?) at higher order it is not a structural tangent but rather a Natural tangent.
But eitherway it follows the requirement that it can be added to its base space to get a nother value in the base space. Its like Real numbers in that way,

oscardssmith · 2024-01-23T14:37:43Z

test/reverse.jl

-
+
+    # while
+    # @test times_three_while'(1.0) == 3.0  # hangs in 1.11


do we know why these hang? that seems bad.

No, but its unrelated to this PR

src/extra_rules.jl

src/stage1/recurse_fwd.jl

src/tangent.jl

…upport

more overloads for StaticArrays

codecov · 2024-02-01T11:44:58Z

Codecov Report

Attention: 16 lines in your changes are missing coverage. Please review.

Comparison is base (6f1bdf0) 55.11% compared to head (8badafd) 55.12%.

Files	Patch %	Lines
src/extra_rules.jl	45.45%	6 Missing ⚠️
src/codegen/forward_demand.jl	40.00%	3 Missing ⚠️
src/stage1/forward.jl	87.50%	3 Missing ⚠️
src/stage1/mixed.jl	0.00%	3 Missing ⚠️
src/higher_fwd_rules.jl	80.00%	1 Missing ⚠️

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #219   +/-   ##
=======================================
  Coverage   55.11%   55.12%           
=======================================
  Files          28       28           
  Lines        2892     2917   +25     
=======================================
+ Hits         1594     1608   +14     
- Misses       1298     1309   +11

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

oxinabox self-assigned this Sep 28, 2023

oxinabox changed the title ~~Forward mode mutation support~~ Forward mode mutable struct support Sep 28, 2023

oxinabox commented Sep 29, 2023

View reviewed changes

test/forward_mutation.jl Outdated Show resolved Hide resolved

oxinabox commented Sep 29, 2023

View reviewed changes

test/forward_mutation.jl Outdated Show resolved Hide resolved

oxinabox force-pushed the ox/mutation2 branch from 7e3c80b to b507348 Compare December 29, 2023 05:24

oxinabox commented Jan 12, 2024

View reviewed changes

src/codegen/forward_demand.jl Outdated Show resolved Hide resolved

oxinabox commented Jan 12, 2024

View reviewed changes

oxinabox commented Jan 16, 2024

View reviewed changes

src/stage1/recurse_fwd.jl Outdated Show resolved Hide resolved

oxinabox force-pushed the ox/mutation2 branch from 7eb435b to 09589cd Compare January 19, 2024 10:47

oxinabox marked this pull request as ready for review January 22, 2024 16:49

oxinabox commented Jan 22, 2024

View reviewed changes

oxinabox force-pushed the ox/mutation2 branch from 09589cd to 0524b08 Compare January 23, 2024 11:48

oscardssmith reviewed Jan 23, 2024

View reviewed changes

oxinabox commented Jan 24, 2024

View reviewed changes

src/extra_rules.jl Outdated Show resolved Hide resolved

oxinabox commented Feb 1, 2024

View reviewed changes

src/extra_rules.jl Outdated Show resolved Hide resolved

oxinabox commented Feb 1, 2024

View reviewed changes

src/stage1/recurse_fwd.jl Outdated Show resolved Hide resolved

oxinabox commented Feb 1, 2024

View reviewed changes

src/tangent.jl Outdated Show resolved Hide resolved

oxinabox added 11 commits February 1, 2024 19:39

WIP: begin switching forward mode over to zero_bundles for mutation s…

8a177b2

…upport

Handle first order constructors

1bffce5

make 2nd order mutation not error

15230b8

fix 2nd order test

8652053

add extra second order test

1d91273

Fix closures

67a15da

remove debug statements

7d481d4

update for more aggressive iszero and take NoTangent more seriously

6ae4e13

type of type erasure in ∂☆new

b742b3e

tests for type of type

e41490f

correct test

d0a825b

oxinabox and others added 12 commits February 1, 2024 19:40

Comment out reverse tests that are making it hang on nigthtly

f47d5f0

fix quotenodes not to error

7c7c628

Correct code for higher order ∂☆new

2b79ef3

more extra rules for static arrays

55eb407

more overloads for StaticArrays

if all partials AbstractZero don't call frule

b9940e3

Remove unrelated confusing comment

d4db011

Remove case in frule for purely AbstractZero tangent

6f57657

Require version of ChainRulesCore with MutableTangent etc

00224ca

Remove outdated comment

a61d05b

Update comment

c9fb55a

Add comment

0cfbc41

Mark tests that broke in 1.11 only broken in 1.11

8badafd

oxinabox force-pushed the ox/mutation2 branch from 9957226 to 8badafd Compare February 1, 2024 11:40

oxinabox merged commit c00ceb7 into main Feb 1, 2024
5 of 8 checks passed

oxinabox deleted the ox/mutation2 branch February 1, 2024 12:58

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Forward mode mutable struct support #219

Forward mode mutable struct support #219

oxinabox commented Sep 28, 2023 •

edited

Loading

oxinabox commented Dec 29, 2023 •

edited

Loading

oxinabox commented Jan 1, 2024

oxinabox Jan 12, 2024

oxinabox Jan 22, 2024

oscardssmith Jan 23, 2024

oxinabox Jan 24, 2024

codecov bot commented Feb 1, 2024

		@@ -362,6 +363,6 @@ function forward_diff!(interp::ADInterpreter, ir::IRCode, src::CodeInfo, mi::Met
		rt = CC._ir_abstract_constant_propagation(interp, irsv)

		ir = compact!(ir)

		# It is a little dubious using StructuralTangent{B} for >1st order, but it is isomorphic.
		# Just watch out for order mixing bugs.



		# while
		# @test times_three_while'(1.0) == 3.0 # hangs in 1.11

Forward mode mutable struct support #219

Forward mode mutable struct support #219

Conversation

oxinabox commented Sep 28, 2023 • edited Loading

oxinabox commented Dec 29, 2023 • edited Loading

oxinabox commented Jan 1, 2024

oxinabox Jan 12, 2024

Choose a reason for hiding this comment

oxinabox Jan 22, 2024

Choose a reason for hiding this comment

oscardssmith Jan 23, 2024

Choose a reason for hiding this comment

oxinabox Jan 24, 2024

Choose a reason for hiding this comment

codecov bot commented Feb 1, 2024

Codecov Report

oxinabox commented Sep 28, 2023 •

edited

Loading

oxinabox commented Dec 29, 2023 •

edited

Loading