366
366
test_rrule (findmax, rand (5 ,3 ))
367
367
@test [0 0 ; 0 5 ] == @inferred unthunk (rrule (findmax, [1 2 ; 3 4 ])[2 ]((5.0 , nothing ))[2 ])
368
368
@test [0 0 ; 0 5 ] == @inferred unthunk (rrule (findmax, [1 2 ; 3 4 ])[2 ]((5.0 , NoTangent ()))[2 ])
369
-
369
+
370
370
# Reverse with dims:
371
371
@test [0 0 ; 5 6 ] == @inferred unthunk (rrule (findmax, [1 2 ; 3 4 ], dims= 1 )[2 ](([5 6 ], nothing ))[2 ])
372
372
@test [5 0 ; 6 0 ] == @inferred unthunk (rrule (findmin, [1 2 ; 3 4 ], dims= 2 )[2 ]((hcat ([5 ,6 ]), nothing ))[2 ])
@@ -379,14 +379,14 @@ end
379
379
# Forward
380
380
test_frule (imum, rand (10 ))
381
381
test_frule (imum, rand (3 ,4 ))
382
- @gpu_broken test_frule (imum, rand (3 ,4 ), fkwargs= (dims= 1 ,))
382
+ @gpu test_frule (imum, rand (3 ,4 ), fkwargs= (dims= 1 ,))
383
383
test_frule (imum, [rand (2 ) for _ in 1 : 3 ])
384
384
test_frule (imum, [rand (2 ) for _ in 1 : 3 , _ in 1 : 4 ]; fkwargs= (dims= 1 ,))
385
385
386
386
# Reverse
387
387
test_rrule (imum, rand (10 ))
388
388
test_rrule (imum, rand (3 ,4 ))
389
- @gpu_broken test_rrule (imum, rand (3 ,4 ), fkwargs= (dims= 1 ,))
389
+ @gpu test_rrule (imum, rand (3 ,4 ), fkwargs= (dims= 1 ,))
390
390
test_rrule (imum, rand (3 ,4 ,5 ), fkwargs= (dims= (1 ,3 ),))
391
391
392
392
# Arrays of arrays
0 commit comments