remove zeros_T_AT and rename cuda ext (#349)

* remove `zeros_T_AT` * Format .jl files (#350) Co-authored-by: prbzrg <prbzrg@users.noreply.github.com> * remove `rand_cstm_AT` --------- Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com> Co-authored-by: prbzrg <prbzrg@users.noreply.github.com>
impICNF · Nov 19, 2023 · 9b4662d · 9b4662d
1 parent e2b720b
commit 9b4662d
Show file tree

Hide file tree

Showing 8 changed files with 43 additions and 66 deletions.
diff --git a/Project.toml b/Project.toml
@@ -39,7 +39,7 @@ Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 
 [extensions]
-CUDAExtICNF = "CUDA"
+ContinuousNormalizingFlowsCUDAExt = "CUDA"
 
 [compat]
 ADTypes = "0.2"

diff --git a/ext/CUDAExtICNF/CUDAExtICNF.jl b/ext/CUDAExtICNF/CUDAExtICNF.jl
diff --git a/ext/CUDAExtICNF/base_cuda.jl b/ext/CUDAExtICNF/base_cuda.jl
diff --git a/ext/ContinuousNormalizingFlowsCUDAExt/ContinuousNormalizingFlowsCUDAExt.jl b/ext/ContinuousNormalizingFlowsCUDAExt/ContinuousNormalizingFlowsCUDAExt.jl
@@ -0,0 +1,10 @@
+module ContinuousNormalizingFlowsCUDAExt
+
+using ContinuousNormalizingFlows, CUDA
+using ContinuousNormalizingFlows.ComputationalResources
+
+@inline function ContinuousNormalizingFlows.rng_AT(::CUDALibs)
+    CURAND.default_rng()
+end
+
+end
diff --git a/src/base.jl b/src/base.jl
@@ -105,23 +105,6 @@ end
     Random.default_rng()
 end
 
-@inline function zeros_T_AT(
-    ::AbstractResource,
-    ::AbstractFlows{T},
-    dims...,
-) where {T <: AbstractFloat}
-    zeros(T, dims...)
-end
-
-@inline function rand_cstm_AT(
-    ::AbstractResource,
-    icnf::AbstractFlows{T},
-    cstm::Any,
-    dims...,
-) where {T <: AbstractFloat}
-    convert.(T, rand(icnf.rng, cstm, dims...))
-end
-
 @views function inference_sol(
     icnf::AbstractFlows{T, <:VectorMode, INPLACE},
     mode::Mode,

diff --git a/src/base_cond_icnf.jl b/src/base_cond_icnf.jl
@@ -10,7 +10,8 @@ export inference, generate, loss
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug_input + n_aug + 1)
+    zrs = similar(xs, n_aug_input + n_aug + 1)
+    @ignore_derivatives fill!(zrs, zero(T))
     ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input)
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
@@ -38,7 +39,8 @@ end
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug_input + n_aug + 1, size(xs, 2))
+    zrs = similar(xs, n_aug_input + n_aug + 1, size(xs, 2))
+    @ignore_derivatives fill!(zrs, zero(T))
     ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input, size(xs, 2))
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
@@ -65,9 +67,10 @@ end
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    new_xs = rand_cstm_AT(icnf.resource, icnf, icnf.basedist)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug + 1)
     ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input)
+    new_xs = oftype(ϵ, rand(icnf.rng, icnf.basedist))
+    zrs = similar(new_xs, n_aug + 1)
+    @ignore_derivatives fill!(zrs, zero(T))
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
             INPLACE,
@@ -94,9 +97,10 @@ end
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    new_xs = rand_cstm_AT(icnf.resource, icnf, icnf.basedist, n)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug + 1, size(new_xs, 2))
-    ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input, size(new_xs, 2))
+    ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input, n)
+    new_xs = oftype(ϵ, rand(icnf.rng, icnf.basedist, n))
+    zrs = similar(new_xs, n_aug + 1, n)
+    @ignore_derivatives fill!(zrs, zero(T))
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
             INPLACE,

diff --git a/src/base_icnf.jl b/src/base_icnf.jl
@@ -9,7 +9,8 @@ export inference, generate, loss
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug_input + n_aug + 1)
+    zrs = similar(xs, n_aug_input + n_aug + 1)
+    @ignore_derivatives fill!(zrs, zero(T))
     ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input)
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
@@ -36,7 +37,8 @@ end
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug_input + n_aug + 1, size(xs, 2))
+    zrs = similar(xs, n_aug_input + n_aug + 1, size(xs, 2))
+    @ignore_derivatives fill!(zrs, zero(T))
     ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input, size(xs, 2))
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
@@ -62,9 +64,10 @@ end
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    new_xs = rand_cstm_AT(icnf.resource, icnf, icnf.basedist)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug + 1)
     ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input)
+    new_xs = oftype(ϵ, rand(icnf.rng, icnf.basedist))
+    zrs = similar(new_xs, n_aug + 1)
+    @ignore_derivatives fill!(zrs, zero(T))
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
             INPLACE,
@@ -90,9 +93,10 @@ end
 ) where {T <: AbstractFloat, INPLACE}
     n_aug = n_augment(icnf, mode)
     n_aug_input = n_augment_input(icnf)
-    new_xs = rand_cstm_AT(icnf.resource, icnf, icnf.basedist, n)
-    zrs = zeros_T_AT(icnf.resource, icnf, n_aug + 1, size(new_xs, 2))
-    ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input, size(new_xs, 2))
+    ϵ = randn(icnf.rng, T, icnf.nvars + n_aug_input, n)
+    new_xs = oftype(ϵ, rand(icnf.rng, icnf.basedist, n))
+    zrs = similar(new_xs, n_aug + 1, n)
+    @ignore_derivatives fill!(zrs, zero(T))
     ODEProblem{INPLACE, SciMLBase.FullSpecialize}(
         ifelse(
             INPLACE,

diff --git a/src/utils.jl b/src/utils.jl
@@ -1,10 +1,11 @@
-function jacobian_batched(
+@views function jacobian_batched(
     icnf::AbstractFlows{T, <:SDVecJacMatrixMode},
     f,
     xs::AbstractMatrix{<:Real},
 ) where {T <: AbstractFloat}
     y = f(xs)
-    z = zeros_T_AT(icnf.resource, icnf, size(xs))
+    z = similar(xs)
+    @ignore_derivatives fill!(z, zero(T))
     res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     Jf = VecJac(f, xs; autodiff = icnf.autodiff_backend)
     for i in axes(xs, 1)
@@ -15,13 +16,14 @@ function jacobian_batched(
     y, copy(res)
 end
 
-function jacobian_batched(
+@views function jacobian_batched(
     icnf::AbstractFlows{T, <:SDJacVecMatrixMode},
     f,
     xs::AbstractMatrix{<:Real},
 ) where {T <: AbstractFloat}
     y = f(xs)
-    z = zeros_T_AT(icnf.resource, icnf, size(xs))
+    z = similar(xs)
+    @ignore_derivatives fill!(z, zero(T))
     res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     Jf = JacVec(f, xs; autodiff = icnf.autodiff_backend)
     for i in axes(xs, 1)
@@ -32,13 +34,14 @@ function jacobian_batched(
     y, copy(res)
 end
 
-function jacobian_batched(
-    icnf::AbstractFlows{T, <:ZygoteMatrixMode},
+@views function jacobian_batched(
+    ::AbstractFlows{T, <:ZygoteMatrixMode},
     f,
     xs::AbstractMatrix{<:Real},
 ) where {T <: AbstractFloat}
     y, VJ = Zygote.pullback(f, xs)
-    z = zeros_T_AT(icnf.resource, icnf, size(xs))
+    z = similar(xs)
+    @ignore_derivatives fill!(z, zero(T))
     res = Zygote.Buffer(xs, size(xs, 1), size(xs, 1), size(xs, 2))
     for i in axes(xs, 1)
         @ignore_derivatives z[i, :] .= one(T)