diff --git a/ext/MathOptAIAbstractGPsExt.jl b/ext/MathOptAIAbstractGPsExt.jl
index 7a65d7e..df17cf1 100644
--- a/ext/MathOptAIAbstractGPsExt.jl
+++ b/ext/MathOptAIAbstractGPsExt.jl
@@ -13,7 +13,7 @@ import MathOptAI
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::MathOptAI.Quantile{<:AbstractGPs.PosteriorGP},
         x::Vector,
     )
@@ -49,7 +49,7 @@ moai_quantile[2] - moai_quantile[1]
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::MathOptAI.Quantile{<:AbstractGPs.PosteriorGP},
     x::Vector,
 )
diff --git a/ext/MathOptAIDecisionTreeExt.jl b/ext/MathOptAIDecisionTreeExt.jl
index 7135c4a..1b73647 100644
--- a/ext/MathOptAIDecisionTreeExt.jl
+++ b/ext/MathOptAIDecisionTreeExt.jl
@@ -12,7 +12,7 @@ import MathOptAI
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::DecisionTree.Root,
         x::Vector,
     )
@@ -49,7 +49,7 @@ julia> y = MathOptAI.add_predictor(model, ml_model, x)
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::DecisionTree.Root,
     x::Vector,
 )
diff --git a/ext/MathOptAIFluxExt.jl b/ext/MathOptAIFluxExt.jl
index dcb875b..a78dcfc 100644
--- a/ext/MathOptAIFluxExt.jl
+++ b/ext/MathOptAIFluxExt.jl
@@ -12,7 +12,7 @@ import MathOptAI
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::Flux.Chain,
         x::Vector;
         config::Dict = Dict{Any,Any}(),
@@ -61,7 +61,7 @@ julia> y = MathOptAI.add_predictor(
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::Flux.Chain,
     x::Vector;
     config::Dict = Dict{Any,Any}(),
diff --git a/ext/MathOptAIGLMExt.jl b/ext/MathOptAIGLMExt.jl
index 9c576e7..cd9409d 100644
--- a/ext/MathOptAIGLMExt.jl
+++ b/ext/MathOptAIGLMExt.jl
@@ -12,7 +12,7 @@ import MathOptAI
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::GLM.LinearModel,
         x::Vector;
         reduced_space::Bool = false,
@@ -39,7 +39,7 @@ julia> y = MathOptAI.add_predictor(model, model_glm, x)
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::GLM.LinearModel,
     x::Vector;
     reduced_space::Bool = false,
@@ -75,7 +75,7 @@ end
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::GLM.GeneralizedLinearModel{
             GLM.GlmResp{Vector{Float64},GLM.Bernoulli{Float64},GLM.LogitLink},
         },
@@ -114,7 +114,7 @@ julia> y = MathOptAI.add_predictor(
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::GLM.GeneralizedLinearModel{
         GLM.GlmResp{Vector{Float64},GLM.Bernoulli{Float64},GLM.LogitLink},
     },
diff --git a/ext/MathOptAILuxExt.jl b/ext/MathOptAILuxExt.jl
index 234724c..b2bc9c2 100644
--- a/ext/MathOptAILuxExt.jl
+++ b/ext/MathOptAILuxExt.jl
@@ -12,7 +12,7 @@ import MathOptAI
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::Tuple{<:Lux.Chain,<:NamedTuple,<:NamedTuple},
         x::Vector;
         config::Dict = Dict{Any,Any}(),
@@ -71,7 +71,7 @@ julia> y = MathOptAI.add_predictor(
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::Tuple{<:Lux.Chain,<:NamedTuple,<:NamedTuple},
     x::Vector;
     config::Dict = Dict{Any,Any}(),
diff --git a/ext/MathOptAIPythonCallExt.jl b/ext/MathOptAIPythonCallExt.jl
index 29e40de..2f28229 100644
--- a/ext/MathOptAIPythonCallExt.jl
+++ b/ext/MathOptAIPythonCallExt.jl
@@ -12,7 +12,7 @@ import MathOptAI
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::MathOptAI.PytorchModel,
         x::Vector;
         config::Dict = Dict{Any,Any}(),
@@ -35,7 +35,7 @@ Add a trained neural network from Pytorch via PythonCall.jl to `model`.
    to control how the activation functions are reformulated.
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::MathOptAI.PytorchModel,
     x::Vector;
     config::Dict = Dict{Any,Any}(),
diff --git a/ext/MathOptAIStatsModelsExt.jl b/ext/MathOptAIStatsModelsExt.jl
index cfbd771..9babe5f 100644
--- a/ext/MathOptAIStatsModelsExt.jl
+++ b/ext/MathOptAIStatsModelsExt.jl
@@ -13,7 +13,7 @@ import StatsModels
 
 """
     MathOptAI.add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::StatsModels.TableRegressionModel,
         x::DataFrames.DataFrame;
         kwargs...,
@@ -62,7 +62,7 @@ julia> test_df.y = MathOptAI.add_predictor(model, predictor, test_df)
 ```
 """
 function MathOptAI.add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::StatsModels.TableRegressionModel,
     df::DataFrames.DataFrame;
     kwargs...,
diff --git a/src/MathOptAI.jl b/src/MathOptAI.jl
index c244cbf..88b73e5 100644
--- a/src/MathOptAI.jl
+++ b/src/MathOptAI.jl
@@ -25,13 +25,12 @@ abstract type AbstractPredictor end
 
 """
     add_predictor(
-        model::JuMP.Model,
+        model::JuMP.AbstractModel,
         predictor::AbstractPredictor,
         x::Vector,
-    )::Vector{JuMP.VariableRef}
+    )::Vector
 
-Return a `Vector{JuMP.VariableRef}` representing `y` such that
-`y = predictor(x)`.
+Return a `Vector` representing `y` such that `y = predictor(x)`.
 
 ## Example
 
@@ -58,7 +57,7 @@ Subject to
 function add_predictor end
 
 """
-    add_predictor(model::JuMP.Model, predictor, x::Matrix)
+    add_predictor(model::JuMP.AbstractModel, predictor, x::Matrix)
 
 Return a `Matrix`, representing `y` such that `y[:, i] = predictor(x[:, i])` for
 each columnn `i`.
@@ -87,7 +86,7 @@ Subject to
  2 x[1,3] + 3 x[2,3] - moai_Affine[1] = 0
 ```
 """
-function add_predictor(model::JuMP.Model, predictor, x::Matrix)
+function add_predictor(model::JuMP.AbstractModel, predictor, x::Matrix)
     y = map(j -> add_predictor(model, predictor, x[:, j]), 1:size(x, 2))
     return reduce(hcat, y)
 end
diff --git a/src/predictors/Affine.jl b/src/predictors/Affine.jl
index 16384d2..59880a7 100644
--- a/src/predictors/Affine.jl
+++ b/src/predictors/Affine.jl
@@ -6,9 +6,9 @@
 
 """
     Affine(
-        A::Matrix{Float64},
-        b::Vector{Float64} = zeros(size(A, 1)),
-    ) <: AbstractPredictor
+        A::Matrix{T},
+        b::Vector{T} = zeros(T, size(A, 1)),
+    ) where {T} <: AbstractPredictor
 
 An [`AbstractPredictor`](@ref) that represents the affine relationship:
 ```math
@@ -41,17 +41,17 @@ julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
  2 x[1] + 3 x[2]
 ```
 """
-struct Affine <: AbstractPredictor
-    A::Matrix{Float64}
-    b::Vector{Float64}
+struct Affine{T} <: AbstractPredictor
+    A::Matrix{T}
+    b::Vector{T}
 end
 
-function Affine(A::Matrix{Float64})
-    return Affine(A, zeros(size(A, 1)))
+function Affine(A::Matrix{T}) where {T}
+    return Affine{T}(A, zeros(T, size(A, 1)))
 end
 
-function Affine(A::Vector{Float64})
-    return Affine(reshape(A, 1, length(A)), [0.0])
+function Affine(A::Vector{T}) where {T}
+    return Affine{T}(reshape(A, 1, length(A)), [zero(T)])
 end
 
 function Base.show(io::IO, p::Affine)
@@ -59,7 +59,7 @@ function Base.show(io::IO, p::Affine)
     return print(io, "Affine(A, b) [input: $n, output: $m]")
 end
 
-function add_predictor(model::JuMP.Model, predictor::Affine, x::Vector)
+function add_predictor(model::JuMP.AbstractModel, predictor::Affine, x::Vector)
     m = size(predictor.A, 1)
     y = JuMP.@variable(model, [1:m], base_name = "moai_Affine")
     bounds = _get_variable_bounds.(x)
@@ -78,8 +78,8 @@ function add_predictor(model::JuMP.Model, predictor::Affine, x::Vector)
 end
 
 function add_predictor(
-    model::JuMP.Model,
-    predictor::ReducedSpace{Affine},
+    model::JuMP.AbstractModel,
+    predictor::ReducedSpace{<:Affine},
     x::Vector,
 )
     A, b = predictor.predictor.A, predictor.predictor.b
diff --git a/src/predictors/BinaryDecisionTree.jl b/src/predictors/BinaryDecisionTree.jl
index 5fb2035..5b93c94 100644
--- a/src/predictors/BinaryDecisionTree.jl
+++ b/src/predictors/BinaryDecisionTree.jl
@@ -69,7 +69,7 @@ function Base.show(io::IO, predictor::BinaryDecisionTree{K,V}) where {K,V}
 end
 
 function add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::BinaryDecisionTree,
     x::Vector;
     atol::Float64 = 0.0,
diff --git a/src/predictors/Pipeline.jl b/src/predictors/Pipeline.jl
index 881d36e..c9fa67a 100644
--- a/src/predictors/Pipeline.jl
+++ b/src/predictors/Pipeline.jl
@@ -59,7 +59,11 @@ function Base.show(io::IO, p::Pipeline)
     return
 end
 
-function add_predictor(model::JuMP.Model, predictor::Pipeline, x::Vector)
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::Pipeline,
+    x::Vector,
+)
     for layer in predictor.layers
         x = add_predictor(model, layer, x)
     end
@@ -67,7 +71,7 @@ function add_predictor(model::JuMP.Model, predictor::Pipeline, x::Vector)
 end
 
 function add_predictor(
-    model::JuMP.Model,
+    model::JuMP.AbstractModel,
     predictor::ReducedSpace{Pipeline},
     x::Vector,
 )
diff --git a/src/predictors/Quantile.jl b/src/predictors/Quantile.jl
index 31b57d5..944ccb6 100644
--- a/src/predictors/Quantile.jl
+++ b/src/predictors/Quantile.jl
@@ -38,7 +38,11 @@ function Base.show(io::IO, q::Quantile)
     return print(io, "Quantile(_, $(q.quantiles))")
 end
 
-function add_predictor(model::JuMP.Model, predictor::Quantile, x::Vector)
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::Quantile,
+    x::Vector,
+)
     M, N = length(x), length(predictor.quantiles)
     y = JuMP.@variable(model, [1:N], base_name = "moai_quantile")
     quantile(q, x...) = Distributions.quantile(predictor.distribution(x...), q)
diff --git a/src/predictors/ReLU.jl b/src/predictors/ReLU.jl
index 7432923..5441660 100644
--- a/src/predictors/ReLU.jl
+++ b/src/predictors/ReLU.jl
@@ -43,15 +43,17 @@ julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
 """
 struct ReLU <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, ::ReLU, x::Vector)
+function add_predictor(model::JuMP.AbstractModel, ::ReLU, x::Vector)
     ub = last.(_get_variable_bounds.(x))
     y = JuMP.@variable(model, [1:length(x)], base_name = "moai_ReLU")
-    _set_bounds_if_finite.(y, 0.0, ub)
+    _set_bounds_if_finite.(y, 0, ub)
     JuMP.@constraint(model, y .== max.(0, x))
     return y
 end
 
-add_predictor(::JuMP.Model, ::ReducedSpace{ReLU}, x::Vector) = max.(0, x)
+function add_predictor(::JuMP.AbstractModel, ::ReducedSpace{ReLU}, x::Vector)
+    return max.(0, x)
+end
 
 """
     ReLUBigM(M::Float64) <: AbstractPredictor
@@ -101,18 +103,22 @@ struct ReLUBigM <: AbstractPredictor
     M::Float64
 end
 
-function add_predictor(model::JuMP.Model, predictor::ReLUBigM, x::Vector)
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::ReLUBigM,
+    x::Vector,
+)
     m = length(x)
     bounds = _get_variable_bounds.(x)
     y = JuMP.@variable(model, [1:m], base_name = "moai_ReLU")
-    _set_bounds_if_finite.(y, 0.0, last.(bounds))
+    _set_bounds_if_finite.(y, 0, last.(bounds))
     for i in 1:m
         lb, ub = bounds[i]
         z = JuMP.@variable(model, binary = true)
         JuMP.@constraint(model, y[i] >= x[i])
         U = min(ub, predictor.M)
         JuMP.@constraint(model, y[i] <= U * z)
-        L = min(max(0.0, -lb), predictor.M)
+        L = min(max(0, -lb), predictor.M)
         JuMP.@constraint(model, y[i] <= x[i] + L * (1 - z))
     end
     return y
@@ -167,13 +173,17 @@ Subject to
 """
 struct ReLUSOS1 <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, predictor::ReLUSOS1, x::Vector)
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::ReLUSOS1,
+    x::Vector,
+)
     m = length(x)
     bounds = _get_variable_bounds.(x)
     y = JuMP.@variable(model, [i in 1:m], base_name = "moai_ReLU")
-    _set_bounds_if_finite.(y, 0.0, last.(bounds))
+    _set_bounds_if_finite.(y, 0, last.(bounds))
     z = JuMP.@variable(model, [1:m], lower_bound = 0, base_name = "_z")
-    _set_bounds_if_finite.(z, -Inf, -first.(bounds))
+    _set_bounds_if_finite.(z, nothing, -first.(bounds))
     JuMP.@constraint(model, x .== y - z)
     for i in 1:m
         JuMP.@constraint(model, [y[i], z[i]] in MOI.SOS1([1.0, 2.0]))
@@ -230,13 +240,17 @@ Subject to
 """
 struct ReLUQuadratic <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, predictor::ReLUQuadratic, x::Vector)
+function add_predictor(
+    model::JuMP.AbstractModel,
+    predictor::ReLUQuadratic,
+    x::Vector,
+)
     m = length(x)
     bounds = _get_variable_bounds.(x)
     y = JuMP.@variable(model, [1:m], base_name = "moai_ReLU")
-    _set_bounds_if_finite.(y, 0.0, last.(bounds))
+    _set_bounds_if_finite.(y, 0, last.(bounds))
     z = JuMP.@variable(model, [1:m], base_name = "_z")
-    _set_bounds_if_finite.(z, 0.0, -first.(bounds))
+    _set_bounds_if_finite.(z, 0, -first.(bounds))
     JuMP.@constraint(model, x .== y - z)
     JuMP.@constraint(model, y .* z .== 0)
     return y
diff --git a/src/predictors/Sigmoid.jl b/src/predictors/Sigmoid.jl
index 3209399..c07e087 100644
--- a/src/predictors/Sigmoid.jl
+++ b/src/predictors/Sigmoid.jl
@@ -45,13 +45,13 @@ julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
 """
 struct Sigmoid <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, ::Sigmoid, x::Vector)
+function add_predictor(model::JuMP.AbstractModel, ::Sigmoid, x::Vector)
     y = JuMP.@variable(model, [1:length(x)], base_name = "moai_Sigmoid")
-    _set_bounds_if_finite.(y, 0.0, 1.0)
+    _set_bounds_if_finite.(y, 0, 1)
     JuMP.@constraint(model, [i in 1:length(x)], y[i] == 1 / (1 + exp(-x[i])))
     return y
 end
 
-function add_predictor(::JuMP.Model, ::ReducedSpace{Sigmoid}, x::Vector)
+function add_predictor(::JuMP.AbstractModel, ::ReducedSpace{Sigmoid}, x::Vector)
     return 1 ./ (1 .+ exp.(-x))
 end
diff --git a/src/predictors/SoftMax.jl b/src/predictors/SoftMax.jl
index 159b525..2fa9e78 100644
--- a/src/predictors/SoftMax.jl
+++ b/src/predictors/SoftMax.jl
@@ -47,19 +47,23 @@ julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
 """
 struct SoftMax <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, ::SoftMax, x::Vector)
+function add_predictor(model::JuMP.AbstractModel, ::SoftMax, x::Vector)
     y = JuMP.@variable(model, [1:length(x)], base_name = "moai_SoftMax")
-    _set_bounds_if_finite.(y, 0.0, 1.0)
+    _set_bounds_if_finite.(y, 0, 1)
     denom = JuMP.@variable(model, base_name = "moai_SoftMax_denom")
-    JuMP.set_lower_bound(denom, 0.0)
+    JuMP.set_lower_bound(denom, 0)
     JuMP.@constraint(model, denom == sum(exp.(x)))
     JuMP.@constraint(model, y .== exp.(x) ./ denom)
     return y
 end
 
-function add_predictor(model::JuMP.Model, ::ReducedSpace{SoftMax}, x::Vector)
+function add_predictor(
+    model::JuMP.AbstractModel,
+    ::ReducedSpace{SoftMax},
+    x::Vector,
+)
     denom = JuMP.@variable(model, base_name = "moai_SoftMax_denom")
-    JuMP.set_lower_bound(denom, 0.0)
+    JuMP.set_lower_bound(denom, 0)
     JuMP.@constraint(model, denom == sum(exp.(x)))
     return exp.(x) ./ denom
 end
diff --git a/src/predictors/SoftPlus.jl b/src/predictors/SoftPlus.jl
index 7917261..1e628f8 100644
--- a/src/predictors/SoftPlus.jl
+++ b/src/predictors/SoftPlus.jl
@@ -43,13 +43,17 @@ julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
 """
 struct SoftPlus <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, ::SoftPlus, x::Vector)
+function add_predictor(model::JuMP.AbstractModel, ::SoftPlus, x::Vector)
     y = JuMP.@variable(model, [1:length(x)], base_name = "moai_SoftPlus")
-    _set_bounds_if_finite.(y, 0.0, Inf)
+    _set_bounds_if_finite.(y, 0, nothing)
     JuMP.@constraint(model, y .== log.(1 .+ exp.(x)))
     return y
 end
 
-function add_predictor(::JuMP.Model, ::ReducedSpace{SoftPlus}, x::Vector)
+function add_predictor(
+    ::JuMP.AbstractModel,
+    ::ReducedSpace{SoftPlus},
+    x::Vector,
+)
     return log.(1 .+ exp.(x))
 end
diff --git a/src/predictors/Tanh.jl b/src/predictors/Tanh.jl
index 2150c01..68c976e 100644
--- a/src/predictors/Tanh.jl
+++ b/src/predictors/Tanh.jl
@@ -45,11 +45,11 @@ julia> y = MathOptAI.add_predictor(model, MathOptAI.ReducedSpace(f), x)
 """
 struct Tanh <: AbstractPredictor end
 
-function add_predictor(model::JuMP.Model, ::Tanh, x::Vector)
+function add_predictor(model::JuMP.AbstractModel, ::Tanh, x::Vector)
     y = JuMP.@variable(model, [1:length(x)], base_name = "moai_Tanh")
-    _set_bounds_if_finite.(y, -1.0, 1.0)
+    _set_bounds_if_finite.(y, -1, 1)
     JuMP.@constraint(model, y .== tanh.(x))
     return y
 end
 
-add_predictor(::JuMP.Model, ::ReducedSpace{Tanh}, x::Vector) = tanh.(x)
+add_predictor(::JuMP.AbstractModel, ::ReducedSpace{Tanh}, x::Vector) = tanh.(x)
diff --git a/src/utilities.jl b/src/utilities.jl
index a814871..8a3055e 100644
--- a/src/utilities.jl
+++ b/src/utilities.jl
@@ -4,8 +4,8 @@
 # Use of this source code is governed by a BSD-style license that can be found
 # in the LICENSE.md file.
 
-function _get_variable_bounds(x::JuMP.VariableRef)
-    lb, ub = -Inf, Inf
+function _get_variable_bounds(x::JuMP.GenericVariableRef{T}) where {T}
+    lb, ub = typemin(T), typemax(T)
     if JuMP.has_upper_bound(x)
         ub = JuMP.upper_bound(x)
     end
@@ -16,16 +16,20 @@ function _get_variable_bounds(x::JuMP.VariableRef)
         lb = ub = JuMP.fix_value(x)
     end
     if JuMP.is_binary(x)
-        lb, ub = max(0.0, lb), min(1.0, ub)
+        lb, ub = max(zero(T), lb), min(one(T), ub)
     end
     return lb, ub
 end
 
-function _set_bounds_if_finite(x::JuMP.VariableRef, l::Float64, u::Float64)
-    if isfinite(l)
+function _set_bounds_if_finite(
+    x::JuMP.GenericVariableRef{T},
+    l::Union{Nothing,Real},
+    u::Union{Nothing,Real},
+) where {T}
+    if l !== nothing && l > typemin(T)
         JuMP.set_lower_bound(x, l)
     end
-    if isfinite(u)
+    if u !== nothing && u < typemax(T)
         JuMP.set_upper_bound(x, u)
     end
     return
@@ -35,4 +39,4 @@ end
 _get_variable_bounds(::Any) = -Inf, Inf
 
 # Default fallback: skip setting variable bound
-_set_bounds_if_finite(::Any, ::Float64, ::Float64) = nothing
+_set_bounds_if_finite(::Any, ::Any, ::Any) = nothing