Debugging and MWEs
There's a reasonable chance that you'll run into an issue with Mooncake.jl at some point. In order to debug what is going on when this happens, or to produce an MWE, it is helpful to have a convenient way to run Mooncake.jl on whatever function and arguments you have which are causing problems.
We recommend making use of Mooncake.jl's testing functionality to generate your test cases:
Mooncake.TestUtils.test_rule
— Functiontest_rule(
rng, x...;
interface_only=false,
is_primitive::Bool=true,
perf_flag::Symbol=:none,
interp::Mooncake.MooncakeInterpreter=Mooncake.get_interpreter(),
debug_mode::Bool=false,
unsafe_perturb::Bool=false,
)
Run standardised tests on the rule
for x
. The first element of x
should be the primal function to test, and each other element a positional argument. In most cases, elements of x
can just be the primal values, and randn_tangent
can be relied upon to generate an appropriate tangent to test. Some notable exceptions exist though, in partcular Ptr
s. In this case, the argument for which randn_tangent
cannot be readily defined should be a CoDual
containing the primal, and a manually constructed tangent field.
This function uses Mooncake.build_rrule
to construct a rule. This will use an rrule!!
if one exists, and derive a rule otherwise.
Arguments
rng::AbstractRNG
: a random number generatorx...
: the function (first element) and its arguments (the remainder)
Keyword Arguments
interface_only::Bool=false
: test only that the interface is satisfied, without testing correctness. This should generally be set tofalse
(the default value), and only enabled if the testing infrastructure is unable to test correctness for some reason e.g. the returned value of the function is aPtr
, and appropriate tangents cannot, therefore, be generated for it automatically.is_primitive::Bool=true
: check whether the thing that you are testing has a hand-writtenrrule!!
. This option is helpful if you are testing a newrrule!!
, as it enables you to verify that your method ofis_primitive
has returned the correct value, and that you are actually testing a method of therrule!!
function – a common mistake when authoring a newrrule!!
is to implementis_primitive
incorrectly and to accidentally wind up testing a rule which Mooncake has derived, as opposed to the one that you have written. If you are testing something for which you have not hand-written anrrule!!
, or which you do not care whether it has a hand-writtenrrule!!
or not, you should set it tofalse
.perf_flag::Symbol=:none
: the value of this symbol determines what kind of performance tests should be performed. By default, none are performed. If you believe that a rule should be allocation-free (iff the primal is allocation free), set this to:allocs
. If you hand-write anrrule!!
and believe that your test case should be type stable, set this to:stability
(at present we cannot verify whether a derived rule is type stable for technical reasons). If you believe that a hand-written rule should be both allocation-free and type-stable, set this to:stability_and_allocs
.interp::Mooncake.MooncakeInterpreter=Mooncake.get_interpreter()
: the abstract interpreter to be used when testing this rule. The default should generally be used.debug_mode::Bool=false
: whether or not the rule should be tested in debug mode. Typically this should be left at its defaultfalse
value, but if you are finding that the tests are failing for a given rule, you may wish to temporarily set it totrue
in order to get access to additional information and automated testing.unsafe_perturb::Bool=false
: value passed as the third argument to_add_to_primal
. Should usually be leftfalse
– consult the docstring for_add_to_primal
for more info on when you might wish to set it totrue
.
This approach is convenient because it can
- check whether AD runs at all,
- check whether AD produces the correct answers,
- check whether AD is performant, and
- can be used without having to manually generate tangents.
Example
For example
f(x) = Core.bitcast(Float64, x)
Mooncake.TestUtils.test_rule(Random.Xoshiro(123), f, 3; is_primitive=false)
will error. (In this particular case, it is caused by Mooncake.jl preventing you from doing (potentially) unsafe casting. In this particular instance, Mooncake.jl just fails to compile, but in other instances other things can happen.)
In any case, the point here is that Mooncake.TestUtils.test_rule
provides a convenient way to produce and report an error.
Segfaults
These are everyone's least favourite kind of problem, and they should be extremely rare in Mooncake.jl. However, if you are unfortunate enough to encounter one, please re-run your problem with the debug_mode
kwarg set to true
. See Debug Mode for more info. In general, this will catch problems before they become segfaults, at which point the above strategy for debugging and error reporting should work well.