Document and finish

Author: oxinabox

src/rule_definition_tools.jl

@@ -135,20 +135,6 @@ function _normalize_scalarrules_macro_input(call, maybe_setup, partials)
     return call, setup_stmts, inputs, partials
 end
 
-"turn both `a` and `a::S` into `a`"
-_unconstrain(arg::Symbol) = arg
-function _unconstrain(arg::Expr)
-    Meta.isexpr(arg, :(::), 2) && return arg.args[1]  # dop constraint.
-    error("malformed arguments: $arg")
-end
-
-"turn both `a` and `::Number` into `a::Number` into `a::Number` etc"
-function _constrain_and_name(arg::Expr, default_constraint)
-    Meta.isexpr(arg, :(::), 2) && return arg  # it is already fine.
-    Meta.isexpr(arg, :(::), 1) && return Expr(:(::), gensym(), arg.args[1])  #add name
-    error("malformed arguments: $arg")
-end
-_constrain_and_name(name::Symbol, constraint) = Expr(:(::), name, constraint)  # add type
 
 function scalar_frule_expr(f, call, setup_stmts, inputs, partials)
     n_outputs = length(partials)
@@ -264,10 +250,47 @@ propagator_name(f::Expr, propname::Symbol) = propagator_name(f.args[end], propna
 propagator_name(fname::Symbol, propname::Symbol) = Symbol(fname, :_, propname)
 propagator_name(fname::QuoteNode, propname::Symbol) = propagator_name(fname.value, propname)
 
+"""
+    @non_differentiable(signature_expression)
+
+A helper to make it easier to declare that a method is not not differentiable.
+This is a short-hand for defining a [`frule`](@ref) and an [`rrule`](@ref)+ pullback that 
+returns [`DoesNotExist()`](@ref) for all partials (except for the function `s̄elf`-partial
+itself which is `NO_FIELDS`)
 
-macro non_differentiable(call_expr)
-    Meta.isexpr(call_expr, :call) || error("Invalid use of `@non_differentiable`")
-    primal_name, orig_args = Iterators.peel(call_expr.args)
+The usage is to put the macro before a function signature.
+Keyword arguments should not be included.
+
+```jldoctest
+julia> @non_differentiable Base.:(==)(a, b)
+
+julia> _, pullback = rrule(==, 2.0, 3.0);
+
+julia> pullback(1.0)
+(Zero(), DoesNotExist(), DoesNotExist())
+```
+
+You can place type-constraints in the signature:
+```jldoctest
+julia> @non_differentiable Base.length(xs::Union{Number, Array})
+
+julia> frule((Zero(), 1), length, [2.0, 3.0])
+(2, DoesNotExist())
+```
+
+!!! warning
+    This helper macro covers only the simple common cases.
+    It does not support Varargs, or `where`-clauses.
+    For these you can declare the `rrule` and `frule` directly
+
+"""
+macro non_differentiable(sig_expr)
+    Meta.isexpr(sig_expr, :call) || error("Invalid use of `@non_differentiable`")
+    for arg in sig_expr.args
+        _isvararg(arg) && error("@non_differentiable does not support Varargs like: $arg")
+    end
+
+    primal_name, orig_args = Iterators.peel(sig_expr.args)
 
     constrained_args = _constrain_and_name.(orig_args, :Any)
     primal_sig_parts = [:(::typeof($primal_name)), constrained_args...]
@@ -304,4 +327,66 @@ function _nondiff_rrule_expr(primal_sig_parts, primal_invoke)
         Expr(:tuple, primal_invoke, pullback_expr),
     )
     return rrule_defn
-end
+end
+
+
+###########
+# Helpers 
+
+"""
+    _isvararg(expr)
+
+returns true if the expression could represent a vararg
+
+```jldoctest
+julia> ChainRulesCore._isvararg(:(x...))
+true
+
+julia> ChainRulesCore._isvararg(:(x::Int...))
+true
+
+julia> ChainRulesCore._isvararg(:(::Int...))
+true
+
+julia> ChainRulesCore._isvararg(:(x::Vararg))
+true
+
+julia> ChainRulesCore._isvararg(:(x::Vararg{Int}))
+true
+
+julia> ChainRulesCore._isvararg(:(::Vararg))
+true
+
+julia> ChainRulesCore._isvararg(:(::Vararg{Int}))
+true
+
+julia> ChainRulesCore._isvararg(:(x))
+false
+````
+"""
+_isvararg(expr) = false
+function _isvararg(expr::Expr)
+    Meta.isexpr(expr, :...) && return true
+    if Meta.isexpr(expr, :(::))
+        constraint = last(expr.args)
+        constraint == :Vararg && return true
+        Meta.isexpr(constraint, :curly) && first(constraint.args) == :Vararg && return true
+    end
+    return false
+end
+
+
+"turn both `a` and `a::S` into `a`"
+_unconstrain(arg::Symbol) = arg
+function _unconstrain(arg::Expr)
+    Meta.isexpr(arg, :(::), 2) && return arg.args[1]  # dop constraint.
+    error("malformed arguments: $arg")
+end
+
+"turn both `a` and `::Number` into `a::Number` into `a::Number` etc"
+function _constrain_and_name(arg::Expr, default_constraint)
+    Meta.isexpr(arg, :(::), 2) && return arg  # it is already fine.
+    Meta.isexpr(arg, :(::), 1) && return Expr(:(::), gensym(), arg.args[1])  #add name
+    error("malformed arguments: $arg")
+end
+_constrain_and_name(name::Symbol, constraint) = Expr(:(::), name, constraint)  # add type

test/rule_definition_tools.jl

@@ -1,5 +1,4 @@
 @testset "rule_definition_tools.jl" begin
-    
     @testset "@non_differentiable" begin
         @testset "nondiff_2_1" begin
             nondiff_2_1(x, y) = fill(7.5, 100)[x + y]
@@ -35,6 +34,32 @@
 
             @test rrule(nonembed_identity, 2.0) == nothing
         end
+
+        @testset "Pointy UnionAll constraints" begin
+            pointy_identity(x) = x
+            @non_differentiable pointy_identity(::Vector{<:AbstractString})
+
+            @test frule((Zero(), 1.2), pointy_identity, ["2"]) == (["2"], DoesNotExist())
+            @test frule((Zero(), 1.2), pointy_identity, 2.0) == nothing
+
+            res, pullback = rrule(pointy_identity, ["2"])
+            @test res == ["2"]
+            @test pullback(1.2) == (NO_FIELDS, DoesNotExist())
+
+            @test rrule(pointy_identity, 2.0) == nothing
+        end
+        
+        @testset "Not supported (Yet)" begin
+            # Varargs are not supported
+            @test_throws Exception @macroexpand(@non_differentiable vararg1(xs...))|
+            @test_throws Exception @macroexpand(@non_differentiable vararg1(xs::Vararg))
+
+            # Where clauses are not supported.
+            @test_throws Exception @macroexpand(
+                @non_differentiable where_identity(::Vector{T}) where T<:AbstractString
+            )
+        end
+
     end
 end