Skip to content

fix function inference #1022

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 12, 2025
Merged

fix function inference #1022

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
4cbee4a
fix isinplace inference and add inference tests
oscardssmith May 12, 2025
de658fd
fixes
oscardssmith May 12, 2025
95f608b
fixes
oscardssmith May 12, 2025
4b195f5
address review
oscardssmith May 12, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions src/problems/optimization_problems.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -131,8 +131,7 @@ function OptimizationProblem(
OptimizationProblem{isinplace(f)}(f, args...; kwargs...)
end
function OptimizationProblem(f, args...; kwargs...)
isinplace(f, 2, has_two_dispatches = false)
OptimizationProblem{true}(OptimizationFunction{true}(f), args...; kwargs...)
OptimizationProblem(OptimizationFunction(f), args...; kwargs...)
end

function OptimizationFunction(
Expand Down
28 changes: 18 additions & 10 deletions src/scimlfunctions.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -948,7 +948,7 @@ dt: the time step

```julia
ImplicitDiscreteFunction{iip,specialize}(f;
analytic = __has_analytic(f) ? f.analytic : nothing,
analytic = __has_analytic(f) ? f.analytic : nothing,
resid_prototype = __has_resid_prototype(f) ? f.resid_prototype : nothing)
```

Expand Down Expand Up @@ -2107,7 +2107,7 @@ A representation of a ODE function `f` with inputs, defined by:
```math
\frac{dx}{dt} = f(x, u, p, t)
```
where `x` are the states of the system and `u` are the inputs (which may represent
where `x` are the states of the system and `u` are the inputs (which may represent
different things in different contexts, such as control variables in optimal control).

Includes all of its related functions, such as the Jacobian of `f`, its gradient
Expand All @@ -2134,7 +2134,7 @@ ODEInputFunction{iip, specialize}(f;
sys = __has_sys(f) ? f.sys : nothing)
```

`f` should be given as `f(x_out,x,u,p,t)` or `out = f(x,u,p,t)`.
`f` should be given as `f(x_out,x,u,p,t)` or `out = f(x,u,p,t)`.
See the section on `iip` for more details on in-place vs out-of-place handling.

- `mass_matrix`: the mass matrix `M` represented in the BVP function. Can be used
Expand Down Expand Up @@ -4199,7 +4199,10 @@ IntervalNonlinearFunction(f::IntervalNonlinearFunction; kwargs...) = f
struct NoAD <: AbstractADType end

(f::OptimizationFunction)(args...) = f.f(args...)
OptimizationFunction(args...; kwargs...) = OptimizationFunction{true}(args...; kwargs...)
function OptimizationFunction(f, args...; kwargs...)
isinplace(f, 2, outofplace_param_number=2)
OptimizationFunction{true}(f, args...; kwargs...)
end

function OptimizationFunction{iip}(f, adtype::AbstractADType = NoAD();
grad = nothing, fg = nothing, hess = nothing, hv = nothing, fgh = nothing,
Expand Down Expand Up @@ -4251,8 +4254,9 @@ end
(f::MultiObjectiveOptimizationFunction)(args...) = f.f(args...)

# Convenience constructor
function MultiObjectiveOptimizationFunction(args...; kwargs...)
MultiObjectiveOptimizationFunction{true}(args...; kwargs...)
function MultiObjectiveOptimizationFunction(f, args...; kwargs...)
isinplace(f, 3)
MultiObjectiveOptimizationFunction{true}(f, args...; kwargs...)
end

# Constructor with keyword arguments
Expand Down Expand Up @@ -4339,15 +4343,17 @@ function BVPFunction{iip, specialize, twopoint}(f, bc;
if iip_f
jac = update_coefficients! #(J,u,p,t)
else
jac = (u, p, t) -> update_coefficients!(deepcopy(jac_prototype), u, p, t)
jac_prototype_copy = deepcopy(jac_prototype)
jac = (u, p, t) -> update_coefficients!(jac_prototype_copy, u, p, t)
end
end

if bcjac === nothing && isa(bcjac_prototype, AbstractSciMLOperator)
if iip_bc
bcjac = update_coefficients! #(J,u,p,t)
else
bcjac = (u, p, t) -> update_coefficients!(deepcopy(bcjac_prototype), u, p, t)
bcjac_prototype_copy = deepcopy(bcjac_prototype)
bcjac = (u, p, t) -> update_coefficients!(bcjac_prototype_copy, u, p, t)
end
end

Expand Down Expand Up @@ -4512,15 +4518,17 @@ function DynamicalBVPFunction{iip, specialize, twopoint}(f, bc;
if iip_f
jac = update_coefficients! #(J,u,p,t)
else
jac = (u, p, t) -> update_coefficients!(deepcopy(jac_prototype), u, p, t)
jac_prototype_copy = deepcopy(jac_prototype)
jac = (u, p, t) -> update_coefficients!(jac_prototype_copy, u, p, t)
end
end

if bcjac === nothing && isa(bcjac_prototype, AbstractSciMLOperator)
if iip_bc
bcjac = update_coefficients! #(J,u,p,t)
else
bcjac = (u, p, t) -> update_coefficients!(deepcopy(bcjac_prototype), u, p, t)
bcjac_prototype_copy = deepcopy(jac_prototype)
bcjac = (u, p, t) -> update_coefficients!(bcjac_prototype_copy, u, p, t)
end
end

Expand Down
2 changes: 1 addition & 1 deletion test/aqua.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,7 @@ end
# for method_ambiguity in ambs
# @show method_ambiguity
# end
@warn "Number of method ambiguities: $(length(ambs))"
!isempty(ambs) &&@warn "Number of method ambiguities: $(length(ambs))"
@test length(ambs) ≤ 8
end

Expand Down
16 changes: 8 additions & 8 deletions test/function_building_error_messages.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,8 +54,8 @@ ofboth(u, p, t) = u
ofboth(du, u, p, t) = du .= u

ODEFunction(ofboth)
ODEFunction{true}(ofboth)
ODEFunction{false}(ofboth)
@inferred ODEFunction{true}(ofboth)
@inferred ODEFunction{false}(ofboth)

jac(u, t) = [1.0]
@test_throws SciMLBase.TooFewArgumentsError ODEFunction(fiip, jac = jac)
Expand Down Expand Up @@ -428,8 +428,8 @@ nfboth(u, p) = u
nfboth(du, u, p) = du .= u

NonlinearFunction(nfboth)
NonlinearFunction{true}(nfboth)
NonlinearFunction{false}(nfboth)
@inferred NonlinearFunction{true}(nfboth)
@inferred NonlinearFunction{false}(nfboth)

njac(u) = [1.0]
@test_throws SciMLBase.TooFewArgumentsError NonlinearFunction(nfiip, jac = njac)
Expand Down Expand Up @@ -520,8 +520,8 @@ bcfboth(u, p, t) = u
bcfboth(du, u, p, t) = du .= u

BVPFunction(bfboth, bcfboth)
BVPFunction{true}(bfboth, bcfboth)
BVPFunction{false}(bfboth, bcfboth)
@inferred BVPFunction{true}(bfboth, bcfboth)
@inferred BVPFunction{false}(bfboth, bcfboth)

bjac(u, t) = [1.0]
bcjac(u, t) = [1.0]
Expand Down Expand Up @@ -663,8 +663,8 @@ dbcfboth(du, u, p, t) = u
dbcfboth(res, du, u, p, t) = res .= du .- u

DynamicalBVPFunction(dbfboth, dbcfboth)
DynamicalBVPFunction{true}(dbfboth, dbcfboth)
DynamicalBVPFunction{false}(dbfboth, dbcfboth)
@inferred DynamicalBVPFunction{true}(dbfboth, dbcfboth)
@inferred DynamicalBVPFunction{false}(dbfboth, dbcfboth)

dbjac(du, u, t) = [1.0]
dbcjac(du, u, t) = [1.0]
Expand Down
Loading