Add `@not_implemented` #335

devmotion · 2021-04-16T19:49:17Z

This is a first draft that addresses #334. It adds an unexported NotImplemented differential that is constructed with an exported @not_implemented helper. If ChainRulesCore.debug_mode() == true, it tracks debugging information such as module and location where the differential is missing and, possibly, additional user-provided information. It throws an error if one tries to perform computations with the differential (I just implemented some of the methods that are defined for other differentials but maybe the list should be shortened or extended).

Examples:

julia> ChainRulesCore.debug_mode() = false # default

julia> zero(ChainRulesCore.@not_implemented())
ERROR: differential not implemented
Stacktrace:
...

julia> zero(ChainRulesCore.@not_implemented("more info"))
ERROR: differential not implemented
Stacktrace:
...

julia> ChainRulesCore.debug_mode() = true

julia> zero(ChainRulesCore.@not_implemented())
ERROR: differential not implemented @ Main #= REPL[18]:1 =#
Stacktrace:
...

julia> zero(ChainRulesCore.@not_implemented("more info"))
ERROR: differential not implemented @ Main #= REPL[20]:1 =#
Info: more info
Stacktrace:
...

I haven't checked yet if/how breaking it is for Zygote.

codecov-commenter · 2021-04-16T19:52:49Z

Codecov Report

Merging #335 (17935fe) into master (8315231) will decrease coverage by 0.53%.
The diff coverage is 85.71%.

@@            Coverage Diff             @@
##           master     #335      +/-   ##
==========================================
- Coverage   90.04%   89.50%   -0.54%     
==========================================
  Files          13       14       +1     
  Lines         472      543      +71     
==========================================
+ Hits          425      486      +61     
- Misses         47       57      +10

Impacted Files	Coverage Δ
src/ChainRulesCore.jl	`100.00% <ø> (ø)`
src/differentials/notimplemented.jl	`67.85% <67.85%> (ø)`
src/differential_arithmetic.jl	`96.72% <97.61%> (+0.47%)`	⬆️
src/differentials/composite.jl	`82.20% <0.00%> (+0.15%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 8315231...17935fe. Read the comment docs.

devmotion · 2021-04-16T20:51:47Z

I just updated the definition of the differential of SpecialFunctions.besseli locally and checked the MWE in FluxML/Zygote.jl#873:

julia> Zygote.gradient(1.0) do x
           besseli(1, x)
       end

It still throws an error since conj is applied in the pullback defined by @scalar_rule (the full computation is conj(NotImplemented(...)) * var"..."). I assume this means that one should rather use the default conj(x::AbstractDifferential) = x instead of throwing an error.

devmotion · 2021-04-16T21:06:17Z

With the reduced implementations, the MWE does not error anymore. When differentiating with respect to the order the NotImplemented is returned, e.g., one obtains

julia> Zygote.gradient(1.0) do x
           besseli(x, 3.0)
       end
(ChainRulesCore.NotImplemented{Nothing, Nothing, Nothing}(nothing, nothing, nothing),)

devmotion · 2021-04-16T21:20:23Z

It seems a bit unfortunate that one has to define linear operations with NotImplemented since otherwise pullbacks throw an error if one wants that computations with it throw an error. The question is, is it important to throw an error for some computations or should just most operations propagate it?

oxinabox · 2021-04-16T21:21:33Z

Nice.
Let's include strong guidance in the docstrings about not using this if the code could be AD'ed through. Only good if the code would error if tried to AD.
And point out it is most useful for if the function is had multiple returns and you have worked out analysically some of but not all of them.

I guess it could also be useful as a short term hack for e.g. Zygote ADing things that it can't handle but that commonly occur not on the actual path of concern (e.g. some function being called and logged).
Since it is more correct than marking those nondifferentiable.

oxinabox · 2021-04-17T00:10:47Z

Re: errors.
The main problem would be that MethodErrors being thrown
And we want to use a custom NotImplementedException that displays the metadata we gave it.
Or that if that wasn't captured displays a .message saying to rerun with debug mode on to identify the source of this.

So we either need to overload linear operators to display that.
Or maybe we get that into into the MethodError by putting it as tyoeparams. Less pretty but gets it across.

I think overloading to propagate is fine. (Another reason why it is like AbsractZero. Maybe we should give a common super type of AbstractZero and this?)
That would mean it would stick around until you tried to add it to a primal e.g. in a optimizer (this is the behavior I wanted for DoesNotExist originally and still kinda do)
And that could be OK since it doesn't matter how much work you do with this, as long as you never add it to a primal, you never use it. You never get a wrong answer

Also it's worth checking if hiding having fields behinds debug mode is worth it.
Does it allocate if we store these, or does Julia do something with compile time literals.

devmotion · 2021-04-17T17:16:25Z

I wonder if one could avoid the definitions of conj etc. by modifying @scalar_rule, e.g., by defining a single function as suggested in #309 for the propagation expression and specializing it for NotImplemented.

mzgubic · 2021-04-19T09:58:39Z

Thanks, that looks great! Just playing around locally a bit, will comment in a couple of hours

mzgubic · 2021-04-19T11:40:53Z

Overall this looks great, thanks for implementing it!

I share your concern about overloading linear operators for NotImplemented. On one hand it is pretty harmless since as @oxinabox points out we never get a wrong result, but it does feel wrong. I will take a look into #309 to see if we can use that instead (my preferred option), but I don't think that should block this PR.

Would it be advantageous to only keep the * operator defined (needed for @scalar_rule to work), and define errors for common others (+, -, / adjoint transpose etc.) so that most users see the nice error message rather than a MethodError?

Just so we don't forget, a few of things to do before we merge:

version bump
tests
documentation (explaining when to use this macro, should probably go in "Writing Good Rules" section)

mzgubic · 2021-04-19T11:41:46Z

src/differentials/notimplemented.jl

+
+# Linear operators
+Base.adjoint(x::NotImplemented) = x
+Base.transpose(x::NotImplemented) = x


Probably don't need these two since conj falls back to the identical AbstractDifferential definition.

mzgubic · 2021-04-19T11:42:12Z

src/differentials/notimplemented.jl

+
+Base.:+(x::NotImplemented, ::Any) = x
+Base.:+(::Any, x::NotImplemented) = x
+Base.:+(x::NotImplemented, ::NotImplemented) = x


Is there a reason we need these? Perhaps we could make it error?

src/differentials/notimplemented.jl

mzgubic · 2021-04-19T13:18:43Z

Ah, one more thing:

julia> @btime  Zygote.gradient(x -> besseli(x, 1.0), 1.0)[1]
  351.905 ns (3 allocations: 48 bytes)
NotImplemented(SpecialFunctions, #= /Users/mzgubic/JuliaEnvs/test/notimplemented/dev/SpecialFunctions/src/chainrules.jl:42 =#, besseli first derivative is not implemented)

julia>  @btime  Zygote.gradient(x -> besseli(x, 1.0), 1.0)[1]
  373.468 ns (3 allocations: 48 bytes)
NotImplemented()

It doesn't seem that the extra info is slowing it down, I would be in favour of using the extra information without debug mode in order to make it more user friendly

devmotion · 2021-04-19T14:20:30Z

Would it be advantageous to only keep the * operator defined (needed for @scalar_rule to work), and define errors for common others (+, -, / adjoint transpose etc.) so that most users see the nice error message rather than a MethodError?

AFAICT the current implementation of @scalar_rule requires us to define at least conj(::NotImplemented) (already provided by the default definition for AbstractDifferential), *(::NotImplemented, ::Any), and muladd(::NotImplemented, ::Any, ::Any), muladd(::Any, ::Any, ::NotImplemented), and muladd(::NotImplemented, ::Any, ::NotImplemented). But I'll check this more carefully.

devmotion · 2021-04-19T14:24:19Z

It doesn't seem that the extra info is slowing it down, I would be in favour of using the extra information without debug mode in order to make it more user friendly

Nice, thanks for looking into this! Would you like me to perform some more benchmarks? It's great if we can provide better error messages without affecting performance 👍

mzgubic · 2021-04-19T14:35:46Z

The other extreme I guess is where one looks at

julia>  @btime  @not_implemented("myerror")
  0.047 ns (0 allocations: 0 bytes)
NotImplemented()

julia> @btime  @not_implemented("myerror")
  2.712 ns (0 allocations: 0 bytes)
NotImplemented(Main, #= REPL[6]:1 =#, myerror)

But since it's such a small absolute number and the linear operators should constant fold it is probably fine? If you can think of more meaningful benchmarks I would love to see them

src/differentials/notimplemented.jl

oxinabox · 2021-04-19T14:50:36Z

I have marked this as assigned to @mzgubic to take over the line for final review.
since it doesn't need two of us commenting on it.
For what its worth i don't have any problem with defining conj or any other linear operator on top of NotImplemented.
Its a bit like missing propagation.
But better since we know where it came from.

Co-authored-by: Lyndon White <[email protected]>

devmotion · 2021-04-19T21:45:07Z

The other extreme I guess is where one looks at
julia>  @btime  @not_implemented("myerror")
  0.047 ns (0 allocations: 0 bytes)
NotImplemented()

julia> @btime  @not_implemented("myerror")
  2.712 ns (0 allocations: 0 bytes)
NotImplemented(Main, #= REPL[6]:1 =#, myerror)
But since it's such a small absolute number and the linear operators should constant fold it is probably fine? If you can think of more meaningful benchmarks I would love to see them

~~The timings are basically the same if one wraps it in a function:~~

julia> g() = ChainRulesCore.@not_implemented("test")
g (generic function with 2 methods)

julia> ChainRulesCore.debug_mode() = false

julia> @btime g()
  2.306 ns (0 allocations: 0 bytes)
NotImplemented(Main, #= REPL[43]:1 =#, test)

julia> ChainRulesCore.debug_mode() = true

julia> @btime g()
  2.307 ns (0 allocations: 0 bytes)
NotImplemented(Main, #= REPL[43]:1 =#, test)

Since I haven't found an example where the debugging information affects performance I remove the check for debug_mode().

Edit: I missed that I had enabled debug mode in both benchmarks, I observe the same difference if performed correctly:

julia> ChainRulesCore.debug_mode() = false

julia> g() = ChainRulesCore.@not_implemented("test")
g (generic function with 2 methods)

julia> @btime g()
  0.018 ns (0 allocations: 0 bytes)
NotImplemented()

julia> ChainRulesCore.debug_mode() = true

julia> g() = ChainRulesCore.@not_implemented("test")
g (generic function with 1 method

julia> @btime g()
  2.306 ns (0 allocations: 0 bytes)
NotImplemented(Main, #= REPL[10]:1 =#, test)

docs/src/writing_good_rules.md

mzgubic · 2021-04-20T09:41:37Z

docs/src/writing_good_rules.md

+
+One can use [`@not_implemented`](@ref) to mark missing differentials. This is helpful if
+the function has multiple inputs or outputs, and one has worked out analytically and
+implemented some but not all differentials.


this is a nitpick, but in markdown we usually have one sentence per line (rather than the character limit)

Ah OK, I'll fix it.

src/differential_arithmetic.jl

test/differentials/notimplemented.jl

mzgubic · 2021-04-20T10:13:52Z

Thank you for this wonderful contribution! I've added a few minor comments, but I think we are good to go. Let me know when you would like it merged.

oxinabox · 2021-04-20T10:29:46Z

Let me know when you would like it merged.

@devmotion has merge rights on the ChainRules project so can merge when ready

Co-authored-by: Miha Zgubic <[email protected]>

mzgubic · 2021-04-20T10:38:11Z

Let me know when you would like it merged.

@devmotion has merge rights on the ChainRules project so can merge when ready

Oh, ok, I assumed not since the PR is from a fork, please go ahead @devmotion

How do i check who has merge rights? It came up a couple of times now that I wish I knew

src/differentials/notimplemented.jl

Co-authored-by: Lyndon White <[email protected]>

devmotion · 2021-04-20T12:20:36Z

@oxinabox @mzgubic Anything else that you would like to be addressed?

Add @not_implemented

8023f98

devmotion marked this pull request as draft April 16, 2021 19:49

devmotion added 3 commits April 16, 2021 22:55

Update documentation

4a64d6c

Define only a few methods

88b1f36

Do not store debugging information if debugging mode is disabled

e5bbda0

devmotion added 7 commits April 16, 2021 23:35

Simplify implementation

3784197

Throw error when converting to a number

b682810

Add more propagations

acb7625

Fix format

088afad

Pretty printing

b24f2fe

Add notes

f3bcefd

Fix admonition

95c2ed2

mzgubic reviewed Apr 19, 2021

View reviewed changes

oxinabox reviewed Apr 19, 2021

View reviewed changes

src/differentials/notimplemented.jl Outdated Show resolved Hide resolved

oxinabox assigned mzgubic Apr 19, 2021

Apply suggestions

fa8f258

Co-authored-by: Lyndon White <[email protected]>

devmotion added 2 commits April 19, 2021 21:13

Allow minimal set of operations and throw exceptions for standard ops

cae24e0

Throw error when calling zero(::NotImplemented)

a6eb951

devmotion added 8 commits April 19, 2021 23:53

Always track debugging information

94b418b

Simplify NotImplemented

236e7f2

Add differential rules

c1c6340

Add missing muladd definitions

0475f41

More tests

225b15f

Add return statement

d4f6138

Add documentation

3fd23d8

Bump version

891efa7

devmotion changed the title ~~[WIP] Add @not_implemented~~ Add @not_implemented Apr 20, 2021

devmotion marked this pull request as ready for review April 20, 2021 07:58

mzgubic approved these changes Apr 20, 2021

View reviewed changes

devmotion and others added 2 commits April 20, 2021 12:31

Apply suggestions from code review

571f3ea

Co-authored-by: Miha Zgubic <[email protected]>

Add missing combination

6e93151

Co-authored-by: Miha Zgubic <[email protected]>

oxinabox reviewed Apr 20, 2021

View reviewed changes

src/differentials/notimplemented.jl Outdated Show resolved Hide resolved

oxinabox reviewed Apr 20, 2021

View reviewed changes

src/differentials/notimplemented.jl Show resolved Hide resolved

devmotion and others added 4 commits April 20, 2021 12:48

Reformat markdown

7bcee3d

Reorder muladd definitions

ee8bb84

Update docstring of NotImplemented

17935fe

Co-authored-by: Lyndon White <[email protected]>

Remove type parameters and make debugging information mandatory

06058a9

oxinabox approved these changes Apr 20, 2021

View reviewed changes

devmotion merged commit 208c85e into JuliaDiff:master Apr 20, 2021

devmotion deleted the dw/notimplemented branch April 20, 2021 12:57

devmotion mentioned this pull request Apr 20, 2021

Add IS_NOT_IMPLEMENTED? #334

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add `@not_implemented` #335

Add `@not_implemented` #335

devmotion commented Apr 16, 2021

codecov-commenter commented Apr 16, 2021 •

edited

Loading

devmotion commented Apr 16, 2021

devmotion commented Apr 16, 2021

devmotion commented Apr 16, 2021

oxinabox commented Apr 16, 2021

oxinabox commented Apr 17, 2021

devmotion commented Apr 17, 2021

mzgubic commented Apr 19, 2021

mzgubic commented Apr 19, 2021 •

edited

Loading

mzgubic Apr 19, 2021

mzgubic Apr 19, 2021

mzgubic commented Apr 19, 2021 •

edited

Loading

devmotion commented Apr 19, 2021

devmotion commented Apr 19, 2021

mzgubic commented Apr 19, 2021 •

edited

Loading

oxinabox commented Apr 19, 2021

devmotion commented Apr 19, 2021 •

edited

Loading

mzgubic Apr 20, 2021

devmotion Apr 20, 2021

devmotion Apr 20, 2021

mzgubic commented Apr 20, 2021

oxinabox commented Apr 20, 2021

mzgubic commented Apr 20, 2021

devmotion commented Apr 20, 2021

Add @not_implemented #335

Add @not_implemented #335

Conversation

devmotion commented Apr 16, 2021

codecov-commenter commented Apr 16, 2021 • edited Loading

Codecov Report

devmotion commented Apr 16, 2021

devmotion commented Apr 16, 2021

devmotion commented Apr 16, 2021

oxinabox commented Apr 16, 2021

oxinabox commented Apr 17, 2021

devmotion commented Apr 17, 2021

mzgubic commented Apr 19, 2021

mzgubic commented Apr 19, 2021 • edited Loading

mzgubic Apr 19, 2021

Choose a reason for hiding this comment

mzgubic Apr 19, 2021

Choose a reason for hiding this comment

mzgubic commented Apr 19, 2021 • edited Loading

devmotion commented Apr 19, 2021

devmotion commented Apr 19, 2021

mzgubic commented Apr 19, 2021 • edited Loading

oxinabox commented Apr 19, 2021

devmotion commented Apr 19, 2021 • edited Loading

mzgubic Apr 20, 2021

Choose a reason for hiding this comment

devmotion Apr 20, 2021

Choose a reason for hiding this comment

devmotion Apr 20, 2021

Choose a reason for hiding this comment

mzgubic commented Apr 20, 2021

oxinabox commented Apr 20, 2021

mzgubic commented Apr 20, 2021

devmotion commented Apr 20, 2021

Add `@not_implemented` #335

Add `@not_implemented` #335

codecov-commenter commented Apr 16, 2021 •

edited

Loading

mzgubic commented Apr 19, 2021 •

edited

Loading

mzgubic commented Apr 19, 2021 •

edited

Loading

mzgubic commented Apr 19, 2021 •

edited

Loading

devmotion commented Apr 19, 2021 •

edited

Loading