fix: replace one with oneunit for basis computation

DifferentiationInterface/CHANGELOG.md

@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Fixed
+
+- Replace `one` with `oneunit` in basis computation ([#826])
+
 ## [0.7.3]
 
 ### Fixed
@@ -62,6 +66,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 [0.6.54]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.53...DifferentiationInterface-v0.6.54
 [0.6.53]: https://github.com/JuliaDiff/DifferentiationInterface.jl/compare/DifferentiationInterface-v0.6.52...DifferentiationInterface-v0.6.53
 
+[#826]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/826
 [#823]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/823
 [#818]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/818
 [#812]: https://github.com/JuliaDiff/DifferentiationInterface.jl/pull/812

DifferentiationInterface/Project.toml

@@ -61,6 +61,7 @@ Aqua = "0.8.12"
 ChainRulesCore = "1.23.0"
 ComponentArrays = "0.15.27"
 DataFrames = "1.7.0"
+Dates = "1"
 DiffResults = "1.1.0"
 Diffractor = "=0.2.6"
 Enzyme = "0.13.39"
@@ -98,6 +99,7 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Diffractor = "9f5e2b26-1114-432f-b630-d3fe2085c51c"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
@@ -130,6 +132,7 @@ test = [
     "Aqua",
     "ComponentArrays",
     "DataFrames",
+    "Dates",
     "ExplicitImports",
     "JET",
     "JLArrays",

DifferentiationInterface/docs/src/explanation/operators.md

@@ -152,4 +152,4 @@ For same-point preparation, the same rules hold with two modifications:
 
 !!! warning
     These rules hold for the majority of backends, but there are some exceptions.
-    The most important exception is [ReverseDiff](@ref) and its taping mechanism, which is sensitive to control flow inside the function.
+    The most important exception is [ReverseDiff.jl](https://github.com/JuliaDiff/ReverseDiff.jl) and its taping mechanism, which is sensitive to control flow inside the function.

DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl

@@ -189,27 +189,27 @@ end
 function DI.value_and_derivative(
     f::F, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    y, ty = DI.value_and_pushforward(f, backend, x, (one(x),), contexts...)
+    y, ty = DI.value_and_pushforward(f, backend, x, (oneunit(x),), contexts...)
     return y, only(ty)
 end
 
 function DI.value_and_derivative!(
     f::F, der, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    y, _ = DI.value_and_pushforward!(f, (der,), backend, x, (one(x),), contexts...)
+    y, _ = DI.value_and_pushforward!(f, (der,), backend, x, (oneunit(x),), contexts...)
     return y, der
 end
 
 function DI.derivative(
     f::F, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    return only(DI.pushforward(f, backend, x, (one(x),), contexts...))
+    return only(DI.pushforward(f, backend, x, (oneunit(x),), contexts...))
 end
 
 function DI.derivative!(
     f::F, der, backend::AutoForwardDiff, x, contexts::Vararg{DI.Context,C}
 ) where {F,C}
-    DI.pushforward!(f, (der,), backend, x, (one(x),), contexts...)
+    DI.pushforward!(f, (der,), backend, x, (oneunit(x),), contexts...)
     return der
 end
 
@@ -220,7 +220,7 @@ function DI.prepare_derivative_nokwarg(
 ) where {F,C}
     _sig = DI.signature(f, backend, x, contexts...; strict)
     pushforward_prep = DI.prepare_pushforward_nokwarg(
-        strict, f, backend, x, (one(x),), contexts...
+        strict, f, backend, x, (oneunit(x),), contexts...
     )
     return ForwardDiffOneArgDerivativePrep(_sig, pushforward_prep)
 end
@@ -234,7 +234,7 @@ function DI.value_and_derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     y, ty = DI.value_and_pushforward(
-        f, prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, only(ty)
 end
@@ -249,7 +249,7 @@ function DI.value_and_derivative!(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     y, _ = DI.value_and_pushforward!(
-        f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, der
 end
@@ -263,7 +263,7 @@ function DI.derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     return only(
-        DI.pushforward(f, prep.pushforward_prep, backend, x, (one(x),), contexts...)
+        DI.pushforward(f, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...)
     )
 end
 
@@ -276,7 +276,9 @@ function DI.derivative!(
     contexts::Vararg{DI.Context,C},
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
-    DI.pushforward!(f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...)
+    DI.pushforward!(
+        f, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
+    )
     return der
 end
 
@@ -638,9 +640,9 @@ function DI.second_derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     return myderivative(T, myderivative(T2, ydual))
@@ -656,9 +658,9 @@ function DI.second_derivative!(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     return myderivative!(T, der2, myderivative(T2, ydual))
@@ -673,9 +675,9 @@ function DI.value_derivative_and_second_derivative(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     y = myvalue(T, myvalue(T2, ydual))
@@ -695,9 +697,9 @@ function DI.value_derivative_and_second_derivative!(
 ) where {F,C}
     DI.check_prep(f, prep, backend, x, contexts...)
     T = tag_type(f, backend, x)
-    xdual = make_dual(T, x, one(x))
+    xdual = make_dual(T, x, oneunit(x))
     T2 = tag_type(f, backend, xdual)
-    xdual2 = make_dual(T2, xdual, one(xdual))
+    xdual2 = make_dual(T2, xdual, oneunit(xdual))
     contexts_dual = translate(typeof(xdual2), contexts)
     ydual = f(xdual2, contexts_dual...)
     y = myvalue(T, myvalue(T2, ydual))
@@ -756,7 +758,7 @@ function DI.value_gradient_and_hessian!(
         contexts isa NTuple{C,DI.GeneralizedConstant}
     )
         fc = DI.fix_tail(f, map(DI.unwrap, contexts)...)
-        result = DiffResult(one(eltype(x)), (grad, hess))
+        result = DiffResult(oneunit(eltype(x)), (grad, hess))
         result = hessian!(result, fc, x)
         y = DR.value(result)
         grad === DR.gradient(result) || copyto!(grad, DR.gradient(result))
@@ -855,7 +857,7 @@ function DI.value_gradient_and_hessian!(
     DI.check_prep(f, prep, backend, x, contexts...)
     contexts_dual = translate_prepared(contexts, prep.contexts_dual)
     fc = DI.fix_tail(f, contexts_dual...)
-    result = DiffResult(one(eltype(x)), (grad, hess))
+    result = DiffResult(oneunit(eltype(x)), (grad, hess))
     CHK = tag_type(backend) === Nothing
     if CHK
         checktag(prep.result_config, f, x)

DifferentiationInterface/ext/DifferentiationInterfaceGPUArraysCoreExt/DifferentiationInterfaceGPUArraysCoreExt.jl

@@ -6,14 +6,14 @@ using GPUArraysCore: @allowscalar, AbstractGPUArray
 function DI.basis(a::AbstractGPUArray{T}, i) where {T}
     b = similar(a)
     fill!(b, zero(T))
-    @allowscalar b[i] = one(T)
+    @allowscalar b[i] = oneunit(T)
     return b
 end
 
 function DI.multibasis(a::AbstractGPUArray{T}, inds) where {T}
     b = similar(a)
     fill!(b, zero(T))
-    view(b, inds) .= one(T)
+    view(b, inds) .= oneunit(T)
     return b
 end
 

DifferentiationInterface/src/first_order/derivative.jl

@@ -143,7 +143,7 @@ function prepare_derivative_nokwarg(
 ) where {F,C}
     _sig = signature(f, backend, x, contexts...; strict)
     pushforward_prep = prepare_pushforward_nokwarg(
-        strict, f, backend, x, (one(x),), contexts...
+        strict, f, backend, x, (oneunit(x),), contexts...
     )
     return PushforwardDerivativePrep(_sig, pushforward_prep)
 end
@@ -153,7 +153,7 @@ function prepare_derivative_nokwarg(
 ) where {F,C}
     _sig = signature(f!, y, backend, x, contexts...; strict)
     pushforward_prep = prepare_pushforward_nokwarg(
-        strict, f!, y, backend, x, (one(x),), contexts...
+        strict, f!, y, backend, x, (oneunit(x),), contexts...
     )
     return PushforwardDerivativePrep(_sig, pushforward_prep)
 end
@@ -169,7 +169,7 @@ function value_and_derivative(
 ) where {F,C}
     check_prep(f, prep, backend, x, contexts...)
     y, ty = value_and_pushforward(
-        f, prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, only(ty)
 end
@@ -184,7 +184,7 @@ function value_and_derivative!(
 ) where {F,C}
     check_prep(f, prep, backend, x, contexts...)
     y, _ = value_and_pushforward!(
-        f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, der
 end
@@ -197,7 +197,7 @@ function derivative(
     contexts::Vararg{Context,C},
 ) where {F,C}
     check_prep(f, prep, backend, x, contexts...)
-    ty = pushforward(f, prep.pushforward_prep, backend, x, (one(x),), contexts...)
+    ty = pushforward(f, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...)
     return only(ty)
 end
 
@@ -210,7 +210,7 @@ function derivative!(
     contexts::Vararg{Context,C},
 ) where {F,C}
     check_prep(f, prep, backend, x, contexts...)
-    pushforward!(f, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...)
+    pushforward!(f, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...)
     return der
 end
 
@@ -226,7 +226,7 @@ function value_and_derivative(
 ) where {F,C}
     check_prep(f!, y, prep, backend, x, contexts...)
     y, ty = value_and_pushforward(
-        f!, y, prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f!, y, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, only(ty)
 end
@@ -242,7 +242,7 @@ function value_and_derivative!(
 ) where {F,C}
     check_prep(f!, y, prep, backend, x, contexts...)
     y, _ = value_and_pushforward!(
-        f!, y, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...
+        f!, y, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
     )
     return y, der
 end
@@ -256,7 +256,7 @@ function derivative(
     contexts::Vararg{Context,C},
 ) where {F,C}
     check_prep(f!, y, prep, backend, x, contexts...)
-    ty = pushforward(f!, y, prep.pushforward_prep, backend, x, (one(x),), contexts...)
+    ty = pushforward(f!, y, prep.pushforward_prep, backend, x, (oneunit(x),), contexts...)
     return only(ty)
 end
 
@@ -270,7 +270,9 @@ function derivative!(
     contexts::Vararg{Context,C},
 ) where {F,C}
     check_prep(f!, y, prep, backend, x, contexts...)
-    pushforward!(f!, y, (der,), prep.pushforward_prep, backend, x, (one(x),), contexts...)
+    pushforward!(
+        f!, y, (der,), prep.pushforward_prep, backend, x, (oneunit(x),), contexts...
+    )
     return der
 end
 

DifferentiationInterface/src/first_order/gradient.jl

@@ -91,7 +91,7 @@ function prepare_gradient_nokwarg(
     _sig = signature(f, backend, x, contexts...; strict)
     y = f(x, map(unwrap, contexts)...)  # TODO: replace with output type inference?
     pullback_prep = prepare_pullback_nokwarg(
-        strict, f, backend, x, (one(typeof(y)),), contexts...
+        strict, f, backend, x, (oneunit(typeof(y)),), contexts...
     )
     return PullbackGradientPrep(_sig, y, pullback_prep)
 end
@@ -106,7 +106,9 @@ function value_and_gradient(
     contexts::Vararg{Context,C},
 ) where {F,SIG,Y,C}
     check_prep(f, prep, backend, x, contexts...)
-    y, tx = value_and_pullback(f, prep.pullback_prep, backend, x, (one(Y),), contexts...)
+    y, tx = value_and_pullback(
+        f, prep.pullback_prep, backend, x, (oneunit(Y),), contexts...
+    )
     return y, only(tx)
 end
 
@@ -120,7 +122,7 @@ function value_and_gradient!(
 ) where {F,SIG,Y,C}
     check_prep(f, prep, backend, x, contexts...)
     y, _ = value_and_pullback!(
-        f, (grad,), prep.pullback_prep, backend, x, (one(Y),), contexts...
+        f, (grad,), prep.pullback_prep, backend, x, (oneunit(Y),), contexts...
     )
     return y, grad
 end
@@ -133,7 +135,7 @@ function gradient(
     contexts::Vararg{Context,C},
 ) where {F,SIG,Y,C}
     check_prep(f, prep, backend, x, contexts...)
-    tx = pullback(f, prep.pullback_prep, backend, x, (one(Y),), contexts...)
+    tx = pullback(f, prep.pullback_prep, backend, x, (oneunit(Y),), contexts...)
     return only(tx)
 end
 
@@ -146,7 +148,7 @@ function gradient!(
     contexts::Vararg{Context,C},
 ) where {F,SIG,Y,C}
     check_prep(f, prep, backend, x, contexts...)
-    pullback!(f, (grad,), prep.pullback_prep, backend, x, (one(Y),), contexts...)
+    pullback!(f, (grad,), prep.pullback_prep, backend, x, (oneunit(Y),), contexts...)
     return grad
 end