fix: mark replicate as non-differentiable

Author: avik-pal

.buildkite/testing.yml

@@ -7,6 +7,9 @@ steps:
               version: "1"
           - JuliaCI/julia-coverage#v1:
               codecov: true
+              dirs:
+                - src
+                - ext
         command: julia --code-coverage=user --color=yes --project .buildkite/scripts/downstream.jl "{{matrix.repo}}" "CUDA"
         agents:
           queue: "juliagpu"
@@ -26,6 +29,9 @@ steps:
               version: "1"
           - JuliaCI/julia-coverage#v1:
               codecov: true
+              dirs:
+                - src
+                - ext
         command: julia --code-coverage=user --color=yes --project .buildkite/scripts/downstream.jl "{{matrix.repo}}" "AMDGPU"
         agents:
           queue: "juliagpu"

.github/workflows/CI.yml

@@ -50,6 +50,8 @@ jobs:
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src,ext
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info
@@ -100,6 +102,8 @@ jobs:
             exit(0)  # Exit immediately, as a success
           end
       - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src,ext
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info
@@ -129,6 +133,8 @@ jobs:
           RETESTITEMS_NWORKERS: 4
           RETESTITEMS_NWORKER_THREADS: 2
       - uses: julia-actions/julia-processcoverage@v1
+        with:
+          directories: src,ext
       - uses: codecov/codecov-action@v4
         with:
           files: lcov.info

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.1.19"
+version = "0.1.20"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -10,10 +10,20 @@ Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 
+[weakdeps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+
+[extensions]
+LuxCoreChainRulesCoreExt = "ChainRulesCore"
+LuxCoreEnzymeCoreExt = "EnzymeCore"
+
 [compat]
 Aqua = "0.8.4"
+ChainRulesCore = "1.24"
 Compat = "4.15.0"
 DispatchDoctor = "0.4.10"
+EnzymeCore = "0.7.7"
 ExplicitImports = "1.9.0"
 Functors = "0.4.8"
 Optimisers = "0.3"

ext/LuxCoreChainRulesCoreExt.jl

@@ -0,0 +1,9 @@
+module LuxCoreChainRulesCoreExt
+
+using ChainRulesCore: @non_differentiable
+using LuxCore: LuxCore
+using Random: AbstractRNG
+
+@non_differentiable LuxCore.replicate(::AbstractRNG)
+
+end

ext/LuxCoreEnzymeCoreExt.jl

@@ -0,0 +1,9 @@
+module LuxCoreEnzymeCoreExt
+
+using EnzymeCore: EnzymeRules
+using LuxCore: LuxCore
+using Random: AbstractRNG
+
+EnzymeRules.inactive(::typeof(LuxCore.replicate), ::AbstractRNG) = nothing
+
+end