Formatting in extensions

.buildkite/pipeline.yml

@@ -1,11 +1,16 @@
+env:
+ SECRET_CODECOV_TOKEN: "nkcRFVXdaPNAbiI0x3qK/XUG8rWjBc8fU73YEyP35SeS465XORqrIYrHUbHuJTRyeyqNRdsHaBcV1P7TBbKAaTQAjHQ1Q0KYfd0uRMSWpZSCgTBz5AwttAxVfFrX+Ky3PzTi2TfDe0uPFZtFo0Asq6sUEr1on+Oo+j+q6br2NK6CrA5yKKuTX4Q2V/UPOIK4vNXY3+zDTKSNtr+HQOlcVEeRIk/0ZQ78Cjd52flEaVw8GWo/CC4YBzLtcOZgaFdgOTEDNHMr0mw6zLE4Y6nxq4lHVSoraSjxjhkB0pXTZ1c51yHX8Jc+q6HC5s87+2Zq5YtsuQSGao+eMtkTAYwfLw==;U2FsdGVkX18z27J3+gNgxsPNnXA0ad4LvZnXeohTam7/6UPqX5+3BYI0tAiVkCho4vlJyL7dd8JEyNtk9BFXsg=="
+
 steps:
  - label: "Julia v1"
  plugins:
  - JuliaCI/julia#v1:
  version: "1"
  - JuliaCI/julia-test#v1: ~
  - JuliaCI/julia-coverage#v1:
- codecov: true
+ dirs:
+ - src
+ - ext
  agents:
  queue: "juliagpu"
  cuda: "*"

ext/MooncakeCUDAExt.jl

@@ -1,71 +1,72 @@
 module MooncakeCUDAExt
 
- using LinearAlgebra, Random, Mooncake
+using LinearAlgebra, Random, Mooncake
 
- using Base: IEEEFloat
- using CUDA: CuArray, cu
+using Base: IEEEFloat
+using CUDA: CuArray, cu
 
- import Mooncake:
-  MinimalCtx,
-  rrule!!,
-  @is_primitive,
-  tangent_type,
-  zero_tangent,
-  randn_tangent,
-  increment!!,
-  set_to_zero!!,
-  _add_to_primal,
-  _diff,
-  _dot,
-  _scale,
-  TestUtils,
-  CoDual,
-  NoPullback
+import Mooncake:
+ MinimalCtx,
+ rrule!!,
+ @is_primitive,
+ tangent_type,
+ zero_tangent,
+ randn_tangent,
+ increment!!,
+ set_to_zero!!,
+ _add_to_primal,
+ _diff,
+ _dot,
+ _scale,
+ TestUtils,
+ CoDual,
+ NoPullback
 
- import Mooncake.TestUtils: populate_address_map!, AddressMap, __increment_should_allocate
+import Mooncake.TestUtils: populate_address_map!, AddressMap, __increment_should_allocate
 
- # Tell Mooncake.jl how to handle CuArrays.
+# Tell Mooncake.jl how to handle CuArrays.
 
- tangent_type(::Type{P}) where {P<:CuArray{<:IEEEFloat}} = P
- zero_tangent(x::CuArray{<:IEEEFloat}) = zero(x)
- function randn_tangent(rng::AbstractRNG, x::CuArray{Float32})
- return cu(randn(rng, Float32, size(x)...))
- end
- TestUtils.has_equal_data(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x == y
- increment!!(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x .+= y
- __increment_should_allocate(::Type{<:CuArray{<:IEEEFloat}}) = true
- set_to_zero!!(x::CuArray{<:IEEEFloat}) = x .= 0
- _add_to_primal(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x + y
- _diff(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x - y
- _dot(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = Float64(dot(x, y))
- _scale(x::Float64, y::P) where {T<:IEEEFloat, P<:CuArray{T}} = T(x) * y
- function populate_address_map!(m::AddressMap, p::CuArray, t::CuArray)
- k = pointer_from_objref(p)
- v = pointer_from_objref(t)
- haskey(m, k) && (@assert m[k] == v)
- m[k] = v
- return m
- end
- function Mooncake._verify_fdata_value(p::CuArray, f::CuArray)
- if size(p) != size(f)
- throw(InvalidFDataException("p has size $(size(p)) but f has size $(size(f))"))
- end
- return nothing
+tangent_type(::Type{P}) where {P<:CuArray{<:IEEEFloat}} = P
+zero_tangent(x::CuArray{<:IEEEFloat}) = zero(x)
+function randn_tangent(rng::AbstractRNG, x::CuArray{Float32})
+ return cu(randn(rng, Float32, size(x)...))
+end
+TestUtils.has_equal_data(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x == y
+increment!!(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x .+= y
+__increment_should_allocate(::Type{<:CuArray{<:IEEEFloat}}) = true
+set_to_zero!!(x::CuArray{<:IEEEFloat}) = x .= 0
+_add_to_primal(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x + y
+_diff(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = x - y
+_dot(x::P, y::P) where {P<:CuArray{<:IEEEFloat}} = Float64(dot(x, y))
+_scale(x::Float64, y::P) where {T<:IEEEFloat, P<:CuArray{T}} = T(x) * y
+function populate_address_map!(m::AddressMap, p::CuArray, t::CuArray)
+ k = pointer_from_objref(p)
+ v = pointer_from_objref(t)
+ haskey(m, k) && (@assert m[k] == v)
+ m[k] = v
+ return m
+end
+function Mooncake._verify_fdata_value(p::CuArray, f::CuArray)
+ if size(p) != size(f)
+ throw(InvalidFDataException("p has size $(size(p)) but f has size $(size(f))"))
  end
+ return nothing
+end
 
- # Basic rules for operating on CuArrays.
+# Basic rules for operating on CuArrays.
+
+@is_primitive(
+ MinimalCtx,
+ Tuple{Type{<:CuArray}, UndefInitializer, Vararg{Int, N}} where {N},
+)
+function rrule!!(
+ p::CoDual{Type{P}},
+ init::CoDual{UndefInitializer},
+ dims::CoDual{Int}...
+) where {P<:CuArray{<:Base.IEEEFloat}}
+ _dims = map(primal, dims)
+ y = CoDual(P(undef, _dims), P(undef, _dims))
+ return y, NoPullback(p, init, dims...)
+end
 
- @is_primitive(
- MinimalCtx,
- Tuple{Type{<:CuArray}, UndefInitializer, Vararg{Int, N}} where {N},
- )
- function rrule!!(
- p::CoDual{Type{P}},
- init::CoDual{UndefInitializer},
- dims::CoDual{Int}...
- ) where {P<:CuArray{<:Base.IEEEFloat}}
- _dims = map(primal, dims)
- y = CoDual(P(undef, _dims), P(undef, _dims))
- return y, NoPullback(p, init, dims...)
- end
 end

ext/MooncakeSpecialFunctionsExt.jl

@@ -1,11 +1,12 @@
 module MooncakeSpecialFunctionsExt
 
- using SpecialFunctions, Mooncake
+using SpecialFunctions, Mooncake
 
- import Mooncake: @from_rrule, DefaultCtx
+import Mooncake: @from_rrule, DefaultCtx
+
+@from_rrule DefaultCtx Tuple{typeof(airyai), Float64}
+@from_rrule DefaultCtx Tuple{typeof(airyaix), Float64}
+@from_rrule DefaultCtx Tuple{typeof(erfc), Float64}
+@from_rrule DefaultCtx Tuple{typeof(erfcx), Float64}
 
- @from_rrule DefaultCtx Tuple{typeof(airyai), Float64}
- @from_rrule DefaultCtx Tuple{typeof(airyaix), Float64}
- @from_rrule DefaultCtx Tuple{typeof(erfc), Float64}
- @from_rrule DefaultCtx Tuple{typeof(erfcx), Float64}
 end

test/integration_testing/cuda.jl

@@ -5,13 +5,13 @@ using CUDA
  # Check we can operate on CuArrays.
  test_tangent(
  Xoshiro(123456),
- CuArray{Float32, 2, CUDA.Mem.DeviceBuffer}(undef, 8, 8);
+ CuArray{Float32, 2, CUDA.DeviceMemory}(undef, 8, 8);
  interface_only=false,
  )
 
  # Check we can instantiate a CuArray.
  test_rule(
- sr(123456), CuArray{Float32, 1, CUDA.Mem.DeviceBuffer}, undef, 256;
+ sr(123456), CuArray{Float32, 1, CUDA.DeviceMemory}, undef, 256;
  interface_only=true, is_primitive=true,
  )
 end