Handle undefined type parameters in fcodual_type and implement rule f…

…or fpext (#225)

* Fix undefined type parameter problem

* Add support for fpext

* Bump patch version

* Remove redundant test case

Project.toml

@@ -1,7 +1,7 @@
 name = "Tapir"
 uuid = "07d77754-e150-4737-8c94-cd238a1fb45b"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.2.33"
+version = "0.2.34"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/codual.jl

@@ -35,7 +35,9 @@ function codual_type(::Type{P}) where {P}
  return isconcretetype(P) ? CoDual{P, tangent_type(P)} : CoDual
 end
 
-codual_type(::Type{Type{P}}) where {P} = CoDual{Type{P}, NoTangent}
+function codual_type(p::Type{Type{P}}) where {P}
+ return @isdefined(P) ? CoDual{Type{P}, NoTangent} : CoDual{_typeof(p), NoTangent}
+end
 
 struct NoPullback{R<:Tuple}
  r::R
@@ -86,6 +88,8 @@ function fcodual_type(::Type{P}) where {P}
  return isconcretetype(P) ? CoDual{P, fdata_type(tangent_type(P))} : CoDual
 end
 
-fcodual_type(::Type{Type{P}}) where {P} = CoDual{Type{P}, NoFData}
+function fcodual_type(p::Type{Type{P}}) where {P}
+ return @isdefined(P) ? CoDual{Type{P}, NoFData} : CoDual{_typeof(p), NoFData}
+end
 
 zero_rdata(x::CoDual) = zero_rdata(primal(x))

src/rrules/builtins.jl

@@ -241,7 +241,13 @@ function rrule!!(::CoDual{typeof(fma_float)}, x, y, z)
  return CoDual(fma_float(_x, _y, primal(z)), NoFData()), fma_float_pullback!!
 end
 
-# fpext -- maybe interesting
+@intrinsic fpext
+function rrule!!(
+ ::CoDual{typeof(fpext)}, ::CoDual{Type{Pext}}, x::CoDual{P}
+) where {Pext<:IEEEFloat, P<:IEEEFloat}
+ fpext_adjoint!!(dy::Pext) = NoRData(), NoRData(), fptrunc(P, dy)
+ return zero_fcodual(fpext(Pext, primal(x))), fpext_adjoint!!
+end
 
 @inactive_intrinsic fpiseq
 @inactive_intrinsic fptosi
@@ -251,7 +257,7 @@ end
 function rrule!!(
  ::CoDual{typeof(fptrunc)}, ::CoDual{Type{Ptrunc}}, x::CoDual{P}
 ) where {Ptrunc<:IEEEFloat, P<:IEEEFloat}
- fptrunc_adjoint!!(dy) = NoRData(), NoRData(), convert(P, dy)
+ fptrunc_adjoint!!(dy::Ptrunc) = NoRData(), NoRData(), convert(P, dy)
  return zero_fcodual(fptrunc(Ptrunc, primal(x))), fptrunc_adjoint!!
 end
 
@@ -823,7 +829,7 @@ function generate_hand_written_rrule!!_test_cases(rng_ctor, ::Val{:builtins})
  (false, :stability, nothing, IntrinsicsWrappers.flipsign_int, 4, -3),
  (false, :stability, nothing, IntrinsicsWrappers.floor_llvm, 4.1),
  (false, :stability, nothing, IntrinsicsWrappers.fma_float, 5.0, 4.0, 3.0),
- # fpext -- NEEDS IMPLEMENTING AND TESTING
+ (true, :stability_and_allocs, nothing, IntrinsicsWrappers.fpext, Float64, 5f0),
  (false, :stability, nothing, IntrinsicsWrappers.fpiseq, 4.1, 4.0),
  (false, :stability, nothing, IntrinsicsWrappers.fptosi, UInt32, 4.1),
  (false, :stability, nothing, IntrinsicsWrappers.fptoui, Int32, 4.1),

test/codual.jl

@@ -7,6 +7,14 @@
  @test codual_type(Real) == CoDual
  @test codual_type(Any) == CoDual
  @test codual_type(Type{UnitRange{Int}}) == CoDual{Type{UnitRange{Int}}, NoTangent}
+ @test ==(
+ codual_type(Type{Tuple{T}} where {T}),
+ CoDual{Type{Type{Tuple{T}} where {T}}, NoTangent},
+ )
+ @test ==(
+ Tapir.fcodual_type(Type{Tuple{T}} where {T}),
+ CoDual{Type{Type{Tuple{T}} where {T}}, NoFData},
+ )
  @test(==(
  codual_type(Union{Float64, Int}),
  Union{CoDual{Float64, Float64}, CoDual{Int, NoTangent}},