Merge #1232

1232: add Flux.skip() r=DhairyaLGandhi a=Moelf

per #821

### PR Checklist

- [x] Tests are added
- [ ] Entry in NEWS.md
- [x] Documentation, if applicable
- [ ] Final review from `@MikeInnes` or `@dhairyagandhi96` (for API changes).


Co-authored-by: Moelf <jerryling315@gmail.com>
Co-authored-by: Moelf <proton@jling.dev>
Co-authored-by: Jerry Ling <proton@jling.dev>

docs/src/utilities.md

@@ -46,4 +46,5 @@ Flux.destructure
 ```@docs
 Flux.throttle
 Flux.stop
+Flux.skip
 ```

src/Flux.jl

@@ -20,6 +20,7 @@ export Chain, Dense, Maxout, RNN, LSTM, GRU, SamePad, Conv, CrossCor, ConvTransp
 include("optimise/Optimise.jl")
 using .Optimise
 using .Optimise: @epochs
+using .Optimise: skip
 export Descent, ADAM, Momentum, Nesterov, RMSProp,
   ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, OADAM,
   ADAMW, RADAM, InvDecay, ExpDecay, WeightDecay,

src/optimise/Optimise.jl

@@ -5,7 +5,7 @@ using LinearAlgebra
 export train!, update!,
 	Descent, ADAM, Momentum, Nesterov, RMSProp,
 	ADAGrad, AdaMax, ADADelta, AMSGrad, NADAM, ADAMW,RADAM, OADAM,
-	InvDecay, ExpDecay, WeightDecay, stop, Optimiser,
+	InvDecay, ExpDecay, WeightDecay, stop, skip, Optimiser,
 	ClipValue, ClipNorm
 
 include("optimisers.jl")

src/optimise/train.jl

@@ -35,6 +35,26 @@ call(f, xs...) = f(xs...)
 runall(f) = f
 runall(fs::AbstractVector) = () -> foreach(call, fs)
 
+struct SkipException <: Exception end
+
+"""
+    skip()
+
+Call `Flux.skip()` in a callback to indicate when a callback condition is met.
+This will trigger the train loop to skip the current data point and not update with the calculated gradient.
+
+# Examples
+```julia
+cb = function ()
+  loss() > 1e7 && Flux.skip()
+end
+```
+"""
+function skip()
+  throw(SkipException())
+end
+
+
 struct StopException <: Exception end
 
 """
@@ -87,6 +107,8 @@ function train!(loss, ps, data, opt; cb = () -> ())
     catch ex
       if ex isa StopException
         break
+      elseif ex isa SkipException
+        continue
       else
         rethrow(ex)
       end

test/optimise.jl

@@ -46,7 +46,25 @@ end
 @testset "Training Loop" begin
   i = 0
   l = 1
+  Flux.train!(
+              () -> (sleep(0.1); Flux.skip(); i+=1),
+              (),
+              Iterators.repeated((), 10),
+              Descent()
+             )
+
+  @test i==0 #all skipped
+
+  Flux.train!(
+              () -> (sleep(0.1); i==8 && Flux.skip(); i+=1),
+              (),
+              Iterators.repeated((), 10),
+              Descent()
+             )
 
+  @test i==8 #skip after i hit 8
+
+  i = 0
   Flux.train!(() -> (sleep(0.1); i += 1; l),
               (),
               Iterators.repeated((), 100),
@@ -121,4 +139,4 @@ end
     @test all(w̄_value .<= 1)
     w̄_norm = Optimise.apply!(ClipNorm(1.0), w, copy(w̄))
     @test norm(w̄_norm) <= 1
-end
+end