Repository: Lando-L/boax (language: python)

Issue:
Add support for constrained optimisation
Start with the expected improvement acquisition function.


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 04ad0fe..243d119 100644
--- a/README.md
+++ b/README.md
@@ -18,6 +18,7 @@ that is **designed for flexibility**. It comes with a low-level interfaces for:
   * Surrogate Models
 * **Constructing and optimizing acquisition functions** (`boax.optimization`):
   * Acquisition Functions
+  * Constraint Functions
   * Maximizers
   * Samplers
 
@@ -43,8 +44,6 @@ For more details check out the [docs](https://boax.readthedocs.io/en/latest/).
 1. Create a dataset from a noisy sinusoid.
 
 ```python
-from functools import partial
-
 from jax import config
 
 # Double precision is highly recommended.
@@ -57,13 +56,11 @@ from jax import numpy as jnp
 from jax import random
 from jax import scipy
 from jax import value_and_grad
-from jax import vmap
 
 import optax
-import matplotlib.pyplot as plt
 
 from boax.prediction import kernels, means, models
-from boax.optimization import acquisitions, maximizers, samplers
+from boax.optimization import acquisitions, maximizers
 
 bounds = jnp.array([[-3, 3]])
 
@@ -81,12 +78,10 @@ y_train = objective(x_train) + 0.3 * random.normal(noise_key, shape=(10,))
 def prior(amplitude, length_scale, noise):
   return models.gaussian_process(
     means.zero(),
-    kernels.scale(nn.softplus(amplitude), kernels.rbf(nn.softplus(length_scale))),
+    kernels.scaled(kernels.rbf(nn.softplus(length_scale)), nn.softplus(amplitude)),
     nn.softplus(noise),
   )
 
-optimizer = optax.adam(0.01)
-
 def target_log_prob(params):
   mean, cov = prior(**params)(x_train)
   return -scipy.stats.multivariate_normal.logpdf(y_train, mean, cov)
@@ -97,6 +92,8 @@ params = {
   'noise': jnp.array(-5.),
 }
 
+optimizer = optax.adam(0.01)
+
 def train_step(state, iteration):
   loss, grads = value_and_grad(target_log_prob)(state[0])
   updates, opt_state = optimizer.update(grads, state[1])
@@ -117,15 +114,15 @@ surrogate = models.gaussian_process_regression(
   x_train,
   y_train,
   means.zero(),
-  kernels.scale(nn.softplus(next_params['amplitude']), kernels.rbf(nn.softplus(next_params['length_scale']))),
+  kernels.scaled(kernels.rbf(nn.softplus(next_params['length_scale'])), nn.softplus(next_params['amplitude'])),
   nn.softplus(next_params['noise']),
 )
 
-acqf = acquisitions.upper_confidence_bound(2.0, surrogate)
+acqf = acquisitions.upper_confidence_bound(surrogate, beta=2.0)
 maximizer = maximizers.bfgs(bounds, q=1, num_restarts=25, num_raw_samples=500)
 
-init_candidates = maximizer.init(maximizer_key, acqf)
-candidates, values = maximizer.maximize(init_candidates, acqf)
+candidates = maximizer.init(maximizer_key, acqf)
+next_candidates, values = maximizer.maximize(candidates, acqf)
 ```
 
 ## Citing Boax
diff --git a/boax/optimization/acquisitions/__init__.py b/boax/optimization/acquisitions/__init__.py
index c83a478..15b4c46 100644
--- a/boax/optimization/acquisitions/__init__.py
+++ b/boax/optimization/acquisitions/__init__.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""The acquisition sub-package."""
+"""The acquisitions sub-package."""
 
 from .alias import expected_improvement as expected_improvement
 from .alias import log_expected_improvement as log_expected_improvement
@@ -27,3 +27,5 @@
 from .alias import q_upper_confidence_bound as q_upper_confidence_bound
 from .alias import upper_confidence_bound as upper_confidence_bound
 from .base import Acquisition as Acquisition
+from .transformed import constrained as constrained
+from .transformed import log_constrained as log_constrained
diff --git a/boax/optimization/acquisitions/alias.py b/boax/optimization/acquisitions/alias.py
index 20d98a3..9b3d81c 100644
--- a/boax/optimization/acquisitions/alias.py
+++ b/boax/optimization/acquisitions/alias.py
@@ -24,12 +24,13 @@
 from boax.optimization.acquisitions import functions
 from boax.optimization.acquisitions.base import Acquisition
 from boax.prediction.models.base import Model
-from boax.typing import Array, Numeric
-from boax.util import compose, tupled
+from boax.utils.functools import compose, tupled
+from boax.utils.stats import mvn_to_norm, sample_mvn, scale_improvement
+from boax.utils.typing import Array, Numeric
 
 
 def probability_of_improvement(
-  best: Numeric, model: Model[Tuple[Array, Array]]
+  model: Model[Tuple[Array, Array]], best: Numeric
 ) -> Acquisition:
   """
   The Probability of Improvement acquisition function.
@@ -39,12 +40,12 @@ def probability_of_improvement(
   `PI(x) = P(y >= best_f), y ~ f(x)`
 
   Example:
-    >>> acqf = probability_of_improvement(0.2, surrogate)
+    >>> acqf = probability_of_improvement(surrogate, 0.2)
     >>> poi = acqf(xs)
 
   Args:
-    best: The best function value observed so far.
     model: A gaussian process regression surrogate model.
+    best: The best function value observed so far.
 
   Returns:
     The corresponding `Acquisition`.
@@ -53,15 +54,15 @@ def probability_of_improvement(
   return jit(
     compose(
       scipy.stats.norm.cdf,
-      tupled(partial(functions.analytic.scaled_improvement, best=best)),
-      tupled(functions.analytic.analytic),
+      tupled(partial(scale_improvement, best=best)),
+      tupled(mvn_to_norm),
       model,
     )
   )
 
 
 def log_probability_of_improvement(
-  best: Numeric, model: Model[Tuple[Array, Array]]
+  model: Model[Tuple[Array, Array]], best: Numeric
 ) -> Acquisition:
   """
   The Log Probability of Improvement acquisition function.
@@ -71,12 +72,12 @@ def log_probability_of_improvement(
   `logPI(x) = log(P(y >= best_f)), y ~ f(x)`
 
   Example:
-    >>> acqf = log_probability_of_improvement(0.2, surrogate)
+    >>> acqf = log_probability_of_improvement(surrogate, 0.2)
     >>> log_poi = acqf(xs)
 
   Args:
-    best: The best function value observed so far.
     model: A gaussian process regression surrogate model.
+    best: The best function value observed so far.
 
   Returns:
     The corresponding `Acquisition`.
@@ -85,15 +86,15 @@ def log_probability_of_improvement(
   return jit(
     compose(
       scipy.stats.norm.logcdf,
-      tupled(partial(functions.analytic.scaled_improvement, best=best)),
-      tupled(functions.analytic.analytic),
+      tupled(partial(scale_improvement, best=best)),
+      tupled(mvn_to_norm),
       model,
     )
   )
 
 
 def expected_improvement(
-  best: Numeric, model: Model[Tuple[Array, Array]]
+  model: Model[Tuple[Array, Array]], best: Numeric
 ) -> Acquisition:
   """
   The Expected Improvement acquisition function.
@@ -105,12 +106,12 @@ def expected_improvement(
   where the expectation is taken over the value of stochastic function `f` at `x`.
 
   Example:
-    >>> acqf = expected_improvement(0.2, surrogate)
+    >>> acqf = expected_improvement(surrogate, 0.2)
     >>> ei = acqf(xs)
 
   Args:
-    best: The best function value observed so far.
     model: A gaussian process regression surrogate model.
+    best: The best function value observed so far.
 
   Returns:
     The corresponding `Acquisition`.
@@ -119,14 +120,14 @@ def expected_improvement(
   return jit(
     compose(
       tupled(partial(functions.analytic.ei, best=best)),
-      tupled(functions.analytic.analytic),
+      tupled(mvn_to_norm),
       model,
     )
   )
 
 
 def log_expected_improvement(
-  best: Numeric, model: Model[Tuple[Array, Array]]
+  model: Model[Tuple[Array, Array]], best: Numeric
 ) -> Acquisition:
   """
   The Log Expected Improvement acquisition function.
@@ -142,12 +143,12 @@ def log_expected_improvement(
     arXiv preprint arXiv:2310.20708 (2023).
 
   Example:
-    >>> acqf = log_expected_improvement(0.2, surrogate)
+    >>> acqf = log_expected_improvement(surrogate, 0.2)
     >>> log_ei = acqf(xs)
 
   Args:
-    best: The best function value observed so far.
     model: A gaussian process regression surrogate model.
+    best: The best function value observed so far.
 
   Returns:
     The corresponding `Acquisition`.
@@ -156,14 +157,14 @@ def log_expected_improvement(
   return jit(
     compose(
       tupled(partial(functions.analytic.lei, best=best)),
-      tupled(functions.analytic.analytic),
+      tupled(mvn_to_norm),
       model,
     )
   )
 
 
 def upper_confidence_bound(
-  beta: Numeric, model: Model[Tuple[Array, Array]]
+  model: Model[Tuple[Array, Array]], beta: Numeric
 ) -> Acquisition:
   """
   The Upper Confidence Bound (UCB) acquisition function.
@@ -174,12 +175,12 @@ def upper_confidence_bound(
   `UCB(x) = loc(x) + sqrt(beta) * scale(x)`
 
   Example:
-    >>> acqf = upper_confidence_bound(2.0, surrogate)
+    >>> acqf = upper_confidence_bound(surrogate, 2.0)
     >>> ucb = acqf(xs)
 
   Args:
-    beta: The mean and covariance trade-off parameter.
     model: A gaussian process regression surrogate model.
+    beta: The mean and covariance trade-off parameter.
 
   Returns:
     The corresponding `Acquisition`.
@@ -188,7 +189,7 @@ def upper_confidence_bound(
   return jit(
     compose(
       tupled(partial(functions.analytic.ucb, beta=beta)),
-      tupled(functions.analytic.analytic),
+      tupled(mvn_to_norm),
       model,
     )
   )
@@ -214,7 +215,7 @@ def posterior_mean(
   return jit(
     compose(
       itemgetter(0),
-      tupled(functions.analytic.analytic),
+      tupled(mvn_to_norm),
       model,
     )
   )
@@ -240,16 +241,16 @@ def posterior_scale(
   return jit(
     compose(
       itemgetter(1),
-      tupled(functions.analytic.analytic),
+      tupled(mvn_to_norm),
       model,
     )
   )
 
 
 def q_expected_improvement(
-  best: Numeric,
-  base_samples: Array,
   model: Model[Tuple[Array, Array]],
+  base_samples: Array,
+  best: Numeric,
 ) -> Acquisition:
   """
   MC-based batch Expected Improvement acquisition function.
@@ -261,13 +262,13 @@ def q_expected_improvement(
   `qEI(x) = E(max(max y - best_f, 0)), y ~ f(x), x = (x_1,...,x_q)`
 
   Example:
-    >>> acqf = q_expected_improvement(0.2, base_samples, surrogate)
+    >>> acqf = q_expected_improvement(surrogate, base_samples, 0.2)
     >>> qei = acqf(xs)
 
   Args:
-    best: The best function value observed so far.
-    base_samples: A set of samples from standard normal distribution.
     model: A gaussian process regression surrogate model.
+    base_samples: A set of samples from standard normal distribution.
+    best: The best function value observed so far.
 
   Returns:
     The corresponding `Acquisition`.
@@ -276,17 +277,17 @@ def q_expected_improvement(
   return jit(
     compose(
       partial(functions.monte_carlo.qei, best=best),
-      tupled(partial(functions.monte_carlo.sampler, base_samples=base_samples)),
+      tupled(partial(sample_mvn, base_samples=base_samples)),
       model,
     )
   )
 
 
 def q_probability_of_improvement(
-  best: Numeric,
-  tau: Numeric,
-  base_samples: Array,
   model: Model[Tuple[Array, Array]],
+  base_samples: Array,
+  tau: Numeric,
+  best: Numeric,
 ) -> Acquisition:
   """
   MC-based batch Probability of Improvement acquisition function.
@@ -300,13 +301,14 @@ def q_probability_of_improvement(
   `qPI(x) = P(max y >= best_f), y ~ f(x), x = (x_1,...,x_q)`
 
   Example:
-    >>> acqf = q_probability_of_improvement(0.2, 1.0, base_samples, surrogate)
+    >>> acqf = q_probability_of_improvement(surrogate, base_samples, 1.0, 0.2)
     >>> qpoi = acqf(xs)
 
   Args:
-    best: The best function value observed so far.
-    base_samples: A set of samples from standard normal distribution.
     model: A gaussian process regression surrogate model.
+    base_samples: A set of samples from standard normal distribution.
+    tau: The temperature parameter.
+    best: The best function value observed so far.
 
   Returns:
     The corresponding `Acquisition`.
@@ -315,16 +317,16 @@ def q_probability_of_improvement(
   return jit(
     compose(
       partial(functions.monte_carlo.qpoi, best=best, tau=tau),
-      tupled(partial(functions.monte_carlo.sampler, base_samples=base_samples)),
+      tupled(partial(sample_mvn, base_samples=base_samples)),
       model,
     )
   )
 
 
 def q_upper_confidence_bound(
-  beta: Numeric,
-  base_samples: Array,
   model: Model[Tuple[Array, Array]],
+  base_samples: Array,
+  beta: Numeric,
 ) -> Acquisition:
   """
   MC-based batch Upper Confidence Bound acquisition function.
@@ -334,13 +336,13 @@ def q_upper_confidence_bound(
   where `y_tilde ~ N(mean(x), beta * pi/2 * cov(x))` and `f(x) ~ N(mean(x), cov(x)).`
 
   Example:
-    >>> acqf = q_upper_confidence_bound(2.0, base_samples, surrogate)
+    >>> acqf = q_upper_confidence_bound(surrogate, base_samples, 2.0)
     >>> qucb = acqf(xs)
 
   Args:
-    beta: The mean and covariance trade-off parameter.
-    base_samples: A set of samples from standard normal distribution.
     model: A gaussian process regression surrogate model.
+    base_samples: A set of samples from standard normal distribution.
+    beta: The mean and covariance trade-off parameter.
 
   Returns:
     The corresponding `Acquisition`.
@@ -350,7 +352,7 @@ def q_upper_confidence_bound(
   return jit(
     compose(
       partial(functions.monte_carlo.qucb, beta=beta_prime),
-      tupled(partial(functions.monte_carlo.sampler, base_samples=base_samples)),
+      tupled(partial(sample_mvn, base_samples=base_samples)),
       model,
     )
   )
diff --git a/boax/optimization/acquisitions/base.py b/boax/optimization/acquisitions/base.py
index dc360f8..c72c2f6 100644
--- a/boax/optimization/acquisitions/base.py
+++ b/boax/optimization/acquisitions/base.py
@@ -16,7 +16,7 @@
 
 from typing import Protocol
 
-from boax.typing import Array
+from boax.utils.typing import Array
 
 
 class Acquisition(Protocol):
diff --git a/boax/optimization/acquisitions/functions/analytic.py b/boax/optimization/acquisitions/functions/analytic.py
index 1f8a8a8..29683c4 100644
--- a/boax/optimization/acquisitions/functions/analytic.py
+++ b/boax/optimization/acquisitions/functions/analytic.py
@@ -15,12 +15,12 @@
 """Analytic acquisition functions."""
 
 import math
-from typing import Tuple
 
 from jax import numpy as jnp
 from jax import scipy
 
-from boax.typing import Array, Numeric
+from boax.utils.stats import scale_improvement
+from boax.utils.typing import Array, Numeric
 
 log2 = math.log(2)
 inv_sqrt2 = 1 / math.sqrt(2)
@@ -28,16 +28,8 @@
 c2 = math.log(math.pi / 2) / 2
 
 
-def analytic(loc: Array, cov: Array) -> Tuple[Array, Array]:
-  return loc, jnp.sqrt(jnp.diag(cov))
-
-
-def scaled_improvement(loc: Array, scale: Array, best: Numeric) -> Array:
-  return (loc - best) / scale
-
-
 def ei(loc: Array, scale: Array, best: Numeric) -> Array:
-  u = scaled_improvement(loc, scale, best)
+  u = scale_improvement(loc, scale, best)
   return (scipy.stats.norm.pdf(u) + u * scipy.stats.norm.cdf(u)) * scale
 
 
@@ -74,7 +66,7 @@ def logh(x):
 
     return jnp.where(x > -1, log_ei_upper(upper), log_ei_lower(lower))
 
-  u = scaled_improvement(loc, scale, best)
+  u = scale_improvement(loc, scale, best)
   return logh(u) + jnp.log(scale)
 
 
diff --git a/boax/optimization/acquisitions/functions/monte_carlo.py b/boax/optimization/acquisitions/functions/monte_carlo.py
index 7f700f6..f4dfe6b 100644
--- a/boax/optimization/acquisitions/functions/monte_carlo.py
+++ b/boax/optimization/acquisitions/functions/monte_carlo.py
@@ -14,14 +14,10 @@
 
 """Monte Carlo acquisition functions."""
 
-from jax import lax, nn, scipy
+from jax import lax, nn
 from jax import numpy as jnp
 
-from boax.typing import Array, Numeric
-
-
-def sampler(mean: Array, cov: Array, base_samples: Array) -> Array:
-  return mean + jnp.dot(scipy.linalg.cholesky(cov, lower=True), base_samples)
+from boax.utils.typing import Array, Numeric
 
 
 def qei(values: Array, best: Numeric) -> Array:
diff --git a/boax/optimization/acquisitions/transformed.py b/boax/optimization/acquisitions/transformed.py
new file mode 100644
index 0000000..da626c6
--- /dev/null
+++ b/boax/optimization/acquisitions/transformed.py
@@ -0,0 +1,69 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformation functions for kernels."""
+
+from jax import lax
+
+from boax.optimization.acquisitions.base import Acquisition
+from boax.optimization.constraints.base import Constraint
+from boax.utils.functools import combine
+
+
+def constrained(
+  acqusition: Acquisition, *constraints: Constraint
+) -> Acquisition:
+  """
+  Constrains a given acquisition function.
+
+  Example:
+    >>> ei = expected_improvement(surrogate, best=0.2)
+    >>> c1 = less_or_equal(cost, 0.8)
+    >>> c2 = less_or_equal(complexity, 8.0)
+    >>> c3 = greater_or_equal(complexity, 2.0)
+    >>> acqf = constrained(ei, c1, c2, c3)
+
+  Args:
+    acquisition: An acquisition function to be constrained.
+    *constraints: The constraints to apply.
+
+  Returns:
+    The constrained `Acquisition`.
+  """
+
+  return combine(lax.mul, 1.0, acqusition, *constraints)
+
+
+def log_constrained(
+  acqusition: Acquisition, *constraints: Constraint
+) -> Acquisition:
+  """
+  Constrains a given log acquisition function.
+
+  Example:
+    >>> ei = log_expected_improvement(surrogate, best=0.2)
+    >>> c1 = log_less_or_equal(cost, 0.8)
+    >>> c2 = log_less_or_equal(complexity, 8.0)
+    >>> c3 = log_greater_or_equal(complexity, 2.0)
+    >>> acqf = log_constrained(ei, c1, c2, c3)
+
+  Args:
+    acquisition: An acquisition function to be constrained.
+    *constrains: The constrains to apply.
+
+  Returns:
+    The constrained `Acquisition`.
+  """
+
+  return combine(lax.add, 0.0, acqusition, *constraints)
diff --git a/boax/optimization/constraints/__init__.py b/boax/optimization/constraints/__init__.py
new file mode 100644
index 0000000..0093150
--- /dev/null
+++ b/boax/optimization/constraints/__init__.py
@@ -0,0 +1,21 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The constrains sub-package."""
+
+from .alias import greater_or_equal as greater_or_equal
+from .alias import less_or_equal as less_or_equal
+from .alias import log_greater_or_equal as log_greater_or_equal
+from .alias import log_less_or_equal as log_less_or_equal
+from .base import Constraint as Constraint
diff --git a/boax/optimization/constraints/alias.py b/boax/optimization/constraints/alias.py
new file mode 100644
index 0000000..e49eca3
--- /dev/null
+++ b/boax/optimization/constraints/alias.py
@@ -0,0 +1,135 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Alias for constrain functions."""
+
+from functools import partial
+from typing import Tuple
+
+from jax import jit
+
+from boax.optimization.constraints import functions
+from boax.optimization.constraints.base import Constraint
+from boax.prediction.models.base import Model
+from boax.utils.functools import compose, tupled
+from boax.utils.stats import mvn_to_norm
+from boax.utils.typing import Array, Numeric
+
+
+def less_or_equal(
+  model: Model[Tuple[Array, Array]], bound: Numeric
+) -> Constraint:
+  """
+  The Less or Equal unequality constraint.
+
+  Example:
+    >>> constraint = less_or_equal(costs, 0.8)
+    >>> le = constraint(xs)
+
+  Args:
+    model: A gaussian process regression feasibility model.
+    best: The lower bound value.
+
+  Returns:
+    The corresponding `Constraint`.
+  """
+
+  return jit(
+    compose(
+      tupled(partial(functions.unequality.le, x=bound)),
+      tupled(mvn_to_norm),
+      model,
+    )
+  )
+
+
+def log_less_or_equal(
+  model: Model[Tuple[Array, Array]], bound: Numeric
+) -> Constraint:
+  """
+  The Log Less or Equal unequality constraint.
+
+  Example:
+    >>> constraint = log_less_or_equal(costs, 0.8)
+    >>> lle = constraint(xs)
+
+  Args:
+    model: A gaussian process regression feasibility model.
+    best: The lower bound value.
+
+  Returns:
+    The corresponding `Constraint`.
+  """
+
+  return jit(
+    compose(
+      tupled(partial(functions.unequality.lle, x=bound)),
+      tupled(mvn_to_norm),
+      model,
+    )
+  )
+
+
+def greater_or_equal(
+  model: Model[Tuple[Array, Array]], bound: Numeric
+) -> Constraint:
+  """
+  The Greater or Equal unequality constraint.
+
+  Example:
+    >>> constraint = greater_or_equal(accuracy, 0.7)
+    >>> ge = constraint(xs)
+
+  Args:
+    model: A gaussian process regression feasibility model.
+    best: The upper bound value.
+
+  Returns:
+    The corresponding `Constraint`.
+  """
+
+  return jit(
+    compose(
+      tupled(partial(functions.unequality.ge, x=bound)),
+      tupled(mvn_to_norm),
+      model,
+    )
+  )
+
+
+def log_greater_or_equal(
+  model: Model[Tuple[Array, Array]], bound: Numeric
+) -> Constraint:
+  """
+  The Log Greater or Equal unequality constraint.
+
+  Example:
+    >>> constraint = log_greater_or_equal(accuracy, 0.7)
+    >>> ge = constraint(xs)
+
+  Args:
+    model: A gaussian process regression feasibility model.
+    best: The upper bound value.
+
+  Returns:
+    The corresponding `Constraint`.
+  """
+
+  return jit(
+    compose(
+      tupled(partial(functions.unequality.lge, x=bound)),
+      tupled(mvn_to_norm),
+      model,
+    )
+  )
diff --git a/boax/optimization/constraints/base.py b/boax/optimization/constraints/base.py
new file mode 100644
index 0000000..e2e1838
--- /dev/null
+++ b/boax/optimization/constraints/base.py
@@ -0,0 +1,39 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Base interface for constrain functions."""
+
+from typing import Protocol
+
+from boax.utils.typing import Array
+
+
+class Constraint(Protocol):
+  """
+  A callable type for constraint functions.
+
+  A constraint function takes a `n x q x d`-dim candidate set as input
+  and returns an `n`-dim array of feasibility scores.
+  """
+
+  def __call__(self, candidates: Array) -> Array:
+    """
+    Evaluates the constraint function on a set of `candidates`.
+
+    Args:
+      candidates: The `n x q x d`-dim candidate set.
+
+    Returns:
+      The `n`-dim feasibility scores at the given design points.
+    """
diff --git a/boax/optimization/constraints/functions/__init__.py b/boax/optimization/constraints/functions/__init__.py
new file mode 100644
index 0000000..b1bf418
--- /dev/null
+++ b/boax/optimization/constraints/functions/__init__.py
@@ -0,0 +1,17 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The acquisition functions sub-package."""
+
+from . import unequality as unequality
diff --git a/boax/optimization/constraints/functions/unequality.py b/boax/optimization/constraints/functions/unequality.py
new file mode 100644
index 0000000..33c66c3
--- /dev/null
+++ b/boax/optimization/constraints/functions/unequality.py
@@ -0,0 +1,35 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Range constraints."""
+
+from jax import scipy
+
+from boax.utils.typing import Array, Numeric
+
+
+def le(loc: Array, scale: Array, x: Numeric) -> Array:
+  return scipy.stats.norm.sf(x, loc, scale)
+
+
+def lle(loc: Array, scale: Array, x: Numeric) -> Array:
+  return scipy.stats.norm.logsf(x, loc, scale)
+
+
+def ge(loc: Array, scale: Array, x: Numeric) -> Array:
+  return scipy.stats.norm.cdf(x, loc, scale)
+
+
+def lge(loc: Array, scale: Array, x: Numeric) -> Array:
+  return scipy.stats.norm.logcdf(x, loc, scale)
diff --git a/boax/optimization/maximizers/alias.py b/boax/optimization/maximizers/alias.py
index bbc9dfc..c5a8c95 100644
--- a/boax/optimization/maximizers/alias.py
+++ b/boax/optimization/maximizers/alias.py
@@ -18,7 +18,7 @@
 
 from boax.optimization.maximizers import functions
 from boax.optimization.maximizers.base import Maximizer
-from boax.typing import Array, Numeric
+from boax.utils.typing import Array, Numeric
 
 
 def bfgs(
@@ -32,10 +32,10 @@ def bfgs(
   The BFGS acquisition function maximizer.
 
   Example:
-    >>> acqf = upper_confidence_bound(2.0, surrogate)
+    >>> acqf = upper_confidence_bound(surrogate, surrogate)
     >>> maximizer = bfgs(bounds, q=1, num_restarts=25, num_raw_samples=500)
-    >>> init_candidates = maximizer.init(key, acqf)
-    >>> candidates, values = maximizer.maximize(init_candidates, acqf)
+    >>> candiates = maximizer.init(key, acqf)
+    >>> next_candidates, values = maximizer.maximize(candiates, acqf)
 
   Args:
     bounds: The bounds of the search space.
diff --git a/boax/optimization/maximizers/base.py b/boax/optimization/maximizers/base.py
index 185510a..6cf5509 100644
--- a/boax/optimization/maximizers/base.py
+++ b/boax/optimization/maximizers/base.py
@@ -17,7 +17,7 @@
 from typing import NamedTuple, Protocol, Tuple
 
 from boax.optimization.acquisitions.base import Acquisition
-from boax.typing import Array, PRNGKey
+from boax.utils.typing import Array, PRNGKey
 
 
 class MaximizerInitFn(Protocol):
diff --git a/boax/optimization/maximizers/functions/initialization.py b/boax/optimization/maximizers/functions/initialization.py
index 9e7e9ab..4d545ad 100644
--- a/boax/optimization/maximizers/functions/initialization.py
+++ b/boax/optimization/maximizers/functions/initialization.py
@@ -19,7 +19,7 @@
 
 from boax.optimization.acquisitions.base import Acquisition
 from boax.optimization.samplers.alias import halton_uniform
-from boax.typing import Array, Numeric, PRNGKey
+from boax.utils.typing import Array, Numeric, PRNGKey
 
 
 def q_batch_initialization(
diff --git a/boax/optimization/maximizers/functions/scipy.py b/boax/optimization/maximizers/functions/scipy.py
index 00a7e41..c0b5102 100644
--- a/boax/optimization/maximizers/functions/scipy.py
+++ b/boax/optimization/maximizers/functions/scipy.py
@@ -21,8 +21,8 @@
 from jax.scipy import optimize
 
 from boax.optimization.acquisitions.base import Acquisition
-from boax.typing import Array
-from boax.util import compose
+from boax.utils.functools import compose
+from boax.utils.typing import Array
 
 
 def bfgs(
diff --git a/boax/optimization/samplers/alias.py b/boax/optimization/samplers/alias.py
index 8df3052..4772325 100644
--- a/boax/optimization/samplers/alias.py
+++ b/boax/optimization/samplers/alias.py
@@ -21,8 +21,8 @@
 
 from boax.optimization.samplers import functions
 from boax.optimization.samplers.base import Sampler
-from boax.typing import Array
-from boax.util import compose
+from boax.utils.functools import compose
+from boax.utils.typing import Array
 
 
 def halton_uniform(
diff --git a/boax/optimization/samplers/base.py b/boax/optimization/samplers/base.py
index 129c448..696457a 100644
--- a/boax/optimization/samplers/base.py
+++ b/boax/optimization/samplers/base.py
@@ -16,7 +16,7 @@
 
 from typing import Protocol
 
-from boax.typing import Array, PRNGKey
+from boax.utils.typing import Array, PRNGKey
 
 
 class Sampler(Protocol):
diff --git a/boax/optimization/samplers/functions/quasi_random.py b/boax/optimization/samplers/functions/quasi_random.py
index f5037f9..c520b56 100644
--- a/boax/optimization/samplers/functions/quasi_random.py
+++ b/boax/optimization/samplers/functions/quasi_random.py
@@ -20,7 +20,7 @@
 from jax import random, scipy
 
 from boax.optimization.samplers.functions.util import primes_less_than
-from boax.typing import Array, PRNGKey
+from boax.utils.typing import Array, PRNGKey
 
 # The maximum dimension we support. This is limited by the number of primes in the PRIMES array.
 sqrt2 = math.sqrt(2)
diff --git a/boax/prediction/kernels/__init__.py b/boax/prediction/kernels/__init__.py
index 2a2c3e2..503fdae 100644
--- a/boax/prediction/kernels/__init__.py
+++ b/boax/prediction/kernels/__init__.py
@@ -21,5 +21,6 @@
 from .alias import rbf as rbf
 from .base import Kernel as Kernel
 from .base import from_kernel_function as from_kernel_function
-from .transformed import combine as combine
-from .transformed import scale as scale
+from .transformed import product as product
+from .transformed import scaled as scaled
+from .transformed import sum as sum
diff --git a/boax/prediction/kernels/alias.py b/boax/prediction/kernels/alias.py
index 8e1ecd4..96a8ac7 100644
--- a/boax/prediction/kernels/alias.py
+++ b/boax/prediction/kernels/alias.py
@@ -18,7 +18,7 @@
 
 from boax.prediction.kernels import functions
 from boax.prediction.kernels.base import Kernel, from_kernel_function
-from boax.typing import Array, Numeric
+from boax.utils.typing import Array, Numeric
 
 
 def rbf(length_scale: Array) -> Kernel:
diff --git a/boax/prediction/kernels/base.py b/boax/prediction/kernels/base.py
index 1d7200d..d4b6d77 100644
--- a/boax/prediction/kernels/base.py
+++ b/boax/prediction/kernels/base.py
@@ -18,7 +18,7 @@
 
 from jax import jit, vmap
 
-from boax.typing import Array, Numeric
+from boax.utils.typing import Array, Numeric
 
 
 class Kernel(Protocol):
diff --git a/boax/prediction/kernels/functions/matern.py b/boax/prediction/kernels/functions/matern.py
index 3f409eb..4404293 100644
--- a/boax/prediction/kernels/functions/matern.py
+++ b/boax/prediction/kernels/functions/matern.py
@@ -19,7 +19,7 @@
 from jax import numpy as jnp
 
 from boax.prediction.kernels.functions.util import euclidean_distance
-from boax.typing import Numeric
+from boax.utils.typing import Numeric
 
 sqrt_3 = math.sqrt(3)
 sqrt_5 = math.sqrt(5)
diff --git a/boax/prediction/kernels/functions/periodic.py b/boax/prediction/kernels/functions/periodic.py
index 9b6edd6..dc07c3f 100644
--- a/boax/prediction/kernels/functions/periodic.py
+++ b/boax/prediction/kernels/functions/periodic.py
@@ -16,7 +16,7 @@
 
 from jax import numpy as jnp
 
-from boax.typing import Numeric
+from boax.utils.typing import Numeric
 
 
 def periodic(
diff --git a/boax/prediction/kernels/functions/rbf.py b/boax/prediction/kernels/functions/rbf.py
index c4df1f6..6d69c84 100644
--- a/boax/prediction/kernels/functions/rbf.py
+++ b/boax/prediction/kernels/functions/rbf.py
@@ -17,7 +17,7 @@
 from jax import numpy as jnp
 
 from boax.prediction.kernels.functions.util import squared_distance
-from boax.typing import Numeric
+from boax.utils.typing import Numeric
 
 
 def rbf(x: Numeric, y: Numeric, length_scale: Numeric) -> Numeric:
diff --git a/boax/prediction/kernels/functions/util.py b/boax/prediction/kernels/functions/util.py
index ce9533d..dbce730 100644
--- a/boax/prediction/kernels/functions/util.py
+++ b/boax/prediction/kernels/functions/util.py
@@ -16,7 +16,7 @@
 
 from jax import numpy as jnp
 
-from boax.typing import Array
+from boax.utils.typing import Array
 
 
 def squared_distance(x: Array, y: Array) -> Array:
diff --git a/boax/prediction/kernels/transformed.py b/boax/prediction/kernels/transformed.py
index dc10281..51c6772 100644
--- a/boax/prediction/kernels/transformed.py
+++ b/boax/prediction/kernels/transformed.py
@@ -14,59 +14,72 @@
 
 """Transformation functions for kernels."""
 
-from typing import Callable
+from functools import partial
 
-from jax import numpy as jnp
+from jax import lax
 
 from boax.prediction.kernels.base import Kernel
-from boax.typing import Numeric
+from boax.utils.functools import combine, compose
+from boax.utils.typing import Numeric
 
 
-def scale(amplitude: Numeric, inner: Kernel) -> Kernel:
+def scaled(kernel: Kernel, amplitude: Numeric) -> Kernel:
   """
-  Scale kernel.
+  Scales a given kernel.
 
   Computes `k(x, y) = amplitude * inner(x, y)`.
 
   Example:
     >>> inner = rbf(jnp.array([0.2, 3.0]))
-    >>> kernel = scale(3.0, inner)
+    >>> kernel = scale(inner, 3.0)
     >>> Kxx = kernel(xs, xs)
 
   Args:
+    kernel: The kernel to be scaled.
     amplitude: The parameter controlling the maximum of the kernel.
-    inner: The inner kernel.
 
   Returns:
     A scaled `Kernel`.
   """
 
-  def kernel(x, y):
-    return amplitude * inner(x, y)
+  return compose(partial(lax.mul, y=amplitude), kernel)
 
-  return kernel
 
+def sum(*kernels: Kernel) -> Kernel:
+  """
+  Constructs a kernel by summing a sequence of kernels.
+
+  Computes `k(x, y) = k1(x, y) + k2(x, y) + ... + kn(x, y)`.
+
+  Example:
+    >>> kernel = sum(map(rbf, [0.2, 0.3, 0.4]))
+    >>> Kxx = kernel(xs, xs)
+
+  Args:
+    kernels: The sequence of kernels to sum.
 
-def combine(operator: Callable, *kernels: Kernel) -> Kernel:
+  Returns:
+    A sum `Kernel`.
+  """
+
+  return combine(lax.add, 0.0, *kernels)
+
+
+def product(*kernels: Kernel) -> Kernel:
   """
-  Combine kernel.
+  Constructs a kernel by multiplying a sequence of kernels.
 
-  Combines a sequence of kernels using the an `operator`.
+  Computes `k(x, y) = k1(x, y) * k2(x, y) * ... * kn(x, y)`.
 
   Example:
-    >>> inners = list(map(rbf, [0.2, 0.3, 0.4]))
-    >>> kernel = combine(partial(jnp.sum, axis=0), *inners)
+    >>> kernel = product(map(rbf, [0.2, 0.3, 0.4]))
     >>> Kxx = kernel(xs, xs)
 
   Args:
-    operator: The operator used for the combination.
     kernels: The sequence of inner kernels.
 
   Returns:
-    A combined `Kernel`.
+    A product `Kernel`.
   """
 
-  def kernel(x, y):
-    return operator(jnp.stack([k(x, y) for k in kernels]))
-
-  return kernel
+  return combine(lax.mul, 1.0, *kernels)
diff --git a/boax/prediction/means/alias.py b/boax/prediction/means/alias.py
index 565f617..80ca0ff 100644
--- a/boax/prediction/means/alias.py
+++ b/boax/prediction/means/alias.py
@@ -18,8 +18,8 @@
 from jax import numpy as jnp
 
 from boax.prediction.means.base import Mean
-from boax.typing import Array, Numeric
-from boax.util import const
+from boax.utils.functools import const
+from boax.utils.typing import Array, Numeric
 
 
 def zero() -> Mean:
diff --git a/boax/prediction/means/base.py b/boax/prediction/means/base.py
index d7fbd54..58e0d22 100644
--- a/boax/prediction/means/base.py
+++ b/boax/prediction/means/base.py
@@ -16,7 +16,7 @@
 
 from typing import Protocol
 
-from boax.typing import Array
+from boax.utils.typing import Array
 
 
 class Mean(Protocol):
diff --git a/boax/prediction/models/alias.py b/boax/prediction/models/alias.py
index 1355ba8..db36928 100644
--- a/boax/prediction/models/alias.py
+++ b/boax/prediction/models/alias.py
@@ -23,7 +23,7 @@
 from boax.prediction.means.base import Mean
 from boax.prediction.models import functions
 from boax.prediction.models.base import Model
-from boax.typing import Array, Numeric
+from boax.utils.typing import Array, Numeric
 
 
 def gaussian_process(
diff --git a/boax/prediction/models/base.py b/boax/prediction/models/base.py
index 7e81473..bb29d0b 100644
--- a/boax/prediction/models/base.py
+++ b/boax/prediction/models/base.py
@@ -16,7 +16,7 @@
 
 from typing import Generic, Protocol, TypeVar
 
-from boax.typing import Array
+from boax.utils.typing import Array
 
 T = TypeVar('T')
 
diff --git a/boax/prediction/models/functions/gaussian.py b/boax/prediction/models/functions/gaussian.py
index 8f4a97e..1d6d61e 100644
--- a/boax/prediction/models/functions/gaussian.py
+++ b/boax/prediction/models/functions/gaussian.py
@@ -21,7 +21,7 @@
 
 from boax.prediction.kernels.base import Kernel
 from boax.prediction.means.base import Mean
-from boax.typing import Array, Numeric
+from boax.utils.typing import Array, Numeric
 
 
 def prior(
diff --git a/boax/utils/__init__.py b/boax/utils/__init__.py
new file mode 100644
index 0000000..0320003
--- /dev/null
+++ b/boax/utils/__init__.py
@@ -0,0 +1,19 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The optimization module."""
+
+from . import functools as functools
+from . import stats as stats
+from . import typing as typing
diff --git a/boax/util.py b/boax/utils/functools.py
similarity index 74%
rename from boax/util.py
rename to boax/utils/functools.py
index 6ccd0dd..c1cabe5 100644
--- a/boax/util.py
+++ b/boax/utils/functools.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""General util functions."""
+"""The functools sub-package."""
 
 from functools import reduce
 from typing import Callable, TypeVar
@@ -53,10 +53,10 @@ def __fn(*args, **kwargs) -> T:
 
 def compose(*fns: Callable) -> Callable:
   """
-  Composes a sequence of function
+  Composes a sequence of functions
 
   Args:
-    fns: Functions to compose.
+    fns: The functions to compose.
 
   Returns:
     A function composed of the sequential applications of the given functions.
@@ -71,6 +71,27 @@ def __fn(*args, **kwargs):
   return reduce(__reduce_fn, fns)
 
 
+def combine(operator: Callable, initial: T, *fns: Callable) -> Callable:
+  """
+  Combines a sequence of functions by an operator.
+
+  Args:
+    operator: The operator used to combine the functions.
+    fns: The functions to combine.
+
+  Returns:
+    A function combining the outputs of the given functions by the operator.
+  """
+
+  def __fn(*args, **kwargs):
+    def __reduce_fn(state: T, f: Callable) -> T:
+      return operator(state, f(*args, **kwargs))
+
+    return reduce(__reduce_fn, fns, initial)
+
+  return __fn
+
+
 def tupled(f: Callable) -> Callable:
   """
   Transforms function to take a tuple as its argument.
diff --git a/boax/utils/stats.py b/boax/utils/stats.py
new file mode 100644
index 0000000..7394164
--- /dev/null
+++ b/boax/utils/stats.py
@@ -0,0 +1,34 @@
+# Copyright 2023 The Boax Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""The stats sub-package."""
+
+from typing import Tuple
+
+from jax import numpy as jnp
+from jax import scipy
+
+from boax.utils.typing import Array
+
+
+def mvn_to_norm(mean: Array, cov: Array) -> Tuple[Array, Array]:
+  return mean, jnp.sqrt(jnp.diag(cov))
+
+
+def sample_mvn(mean: Array, cov: Array, base_samples: Array) -> Array:
+  return mean + jnp.dot(scipy.linalg.cholesky(cov, lower=True), base_samples)
+
+
+def scale_improvement(loc: Array, scale: Array, best: Array) -> Array:
+  return (loc - best) / scale
diff --git a/boax/typing.py b/boax/utils/typing.py
similarity index 95%
rename from boax/typing.py
rename to boax/utils/typing.py
index 300ef09..5c76e03 100644
--- a/boax/typing.py
+++ b/boax/utils/typing.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Alias for JAX typing."""
+"""The typing sub-package."""
 
 import chex
 
diff --git a/docs/source/api_reference/boax.optimization.rst b/docs/source/api_reference/boax.optimization.rst
index be471b5..af8aaab 100644
--- a/docs/source/api_reference/boax.optimization.rst
+++ b/docs/source/api_reference/boax.optimization.rst
@@ -48,6 +48,44 @@ Monte Carlo Acquisition Functions
   q_probability_of_improvement
 
 
+Transformed Acquisition Functions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. autosummary::
+  :toctree: generated
+
+  constrained
+  log_constrained
+
+
+boax.optimization.constraints
+------------------------------
+
+.. currentmodule:: boax.optimization.constraints
+
+
+Constraint Types
+~~~~~~~~~~~~~~~~
+
+.. autoclass:: Constraint
+
+
+Constraint Functions
+~~~~~~~~~~~~~~~~~~~~
+
+
+Unequality Constraint Functions
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+.. autosummary::
+  :toctree: generated
+
+  less_or_equal
+  log_less_or_equal
+  greater_or_equal
+  log_greater_or_equal
+
+
 boax.optimization.maximizers
 ----------------------------
 
diff --git a/docs/source/api_reference/boax.prediction.rst b/docs/source/api_reference/boax.prediction.rst
index 8f0440e..c722789 100644
--- a/docs/source/api_reference/boax.prediction.rst
+++ b/docs/source/api_reference/boax.prediction.rst
@@ -51,6 +51,17 @@ Periodic Kernels
   periodic
 
 
+Transformed Kernels
+^^^^^^^^^^^^^^^^^^^
+
+.. autosummary::
+  :toctree: generated
+
+  scaled
+  sum
+  product
+
+
 boax.prediction.means
 ----------------------------
 
diff --git a/docs/source/api_reference/generated/boax.optimization.acquisitions.constrained.rst b/docs/source/api_reference/generated/boax.optimization.acquisitions.constrained.rst
new file mode 100644
index 0000000..a50d584
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.optimization.acquisitions.constrained.rst
@@ -0,0 +1,6 @@
+﻿boax.optimization.acquisitions.constrained
+==========================================
+
+.. currentmodule:: boax.optimization.acquisitions
+
+.. autofunction:: constrained
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.optimization.acquisitions.log_constrained.rst b/docs/source/api_reference/generated/boax.optimization.acquisitions.log_constrained.rst
new file mode 100644
index 0000000..9164952
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.optimization.acquisitions.log_constrained.rst
@@ -0,0 +1,6 @@
+﻿boax.optimization.acquisitions.log\_constrained
+===============================================
+
+.. currentmodule:: boax.optimization.acquisitions
+
+.. autofunction:: log_constrained
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.optimization.constraints.greater_or_equal.rst b/docs/source/api_reference/generated/boax.optimization.constraints.greater_or_equal.rst
new file mode 100644
index 0000000..376944b
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.optimization.constraints.greater_or_equal.rst
@@ -0,0 +1,6 @@
+﻿boax.optimization.constraints.greater\_or\_equal
+================================================
+
+.. currentmodule:: boax.optimization.constraints
+
+.. autofunction:: greater_or_equal
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.optimization.constraints.less_or_equal.rst b/docs/source/api_reference/generated/boax.optimization.constraints.less_or_equal.rst
new file mode 100644
index 0000000..759e343
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.optimization.constraints.less_or_equal.rst
@@ -0,0 +1,6 @@
+﻿boax.optimization.constraints.less\_or\_equal
+=============================================
+
+.. currentmodule:: boax.optimization.constraints
+
+.. autofunction:: less_or_equal
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.optimization.constraints.log_greater_or_equal.rst b/docs/source/api_reference/generated/boax.optimization.constraints.log_greater_or_equal.rst
new file mode 100644
index 0000000..febf9f5
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.optimization.constraints.log_greater_or_equal.rst
@@ -0,0 +1,6 @@
+﻿boax.optimization.constraints.log\_greater\_or\_equal
+=====================================================
+
+.. currentmodule:: boax.optimization.constraints
+
+.. autofunction:: log_greater_or_equal
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.optimization.constraints.log_less_or_equal.rst b/docs/source/api_reference/generated/boax.optimization.constraints.log_less_or_equal.rst
new file mode 100644
index 0000000..b634c7e
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.optimization.constraints.log_less_or_equal.rst
@@ -0,0 +1,6 @@
+﻿boax.optimization.constraints.log\_less\_or\_equal
+==================================================
+
+.. currentmodule:: boax.optimization.constraints
+
+.. autofunction:: log_less_or_equal
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.prediction.kernels.product.rst b/docs/source/api_reference/generated/boax.prediction.kernels.product.rst
new file mode 100644
index 0000000..e5d242c
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.prediction.kernels.product.rst
@@ -0,0 +1,6 @@
+﻿boax.prediction.kernels.product
+===============================
+
+.. currentmodule:: boax.prediction.kernels
+
+.. autofunction:: product
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.prediction.kernels.scaled.rst b/docs/source/api_reference/generated/boax.prediction.kernels.scaled.rst
new file mode 100644
index 0000000..5f27e33
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.prediction.kernels.scaled.rst
@@ -0,0 +1,6 @@
+﻿boax.prediction.kernels.scaled
+==============================
+
+.. currentmodule:: boax.prediction.kernels
+
+.. autofunction:: scaled
\ No newline at end of file
diff --git a/docs/source/api_reference/generated/boax.prediction.kernels.sum.rst b/docs/source/api_reference/generated/boax.prediction.kernels.sum.rst
new file mode 100644
index 0000000..a3dfb12
--- /dev/null
+++ b/docs/source/api_reference/generated/boax.prediction.kernels.sum.rst
@@ -0,0 +1,6 @@
+﻿boax.prediction.kernels.sum
+===========================
+
+.. currentmodule:: boax.prediction.kernels
+
+.. autofunction:: sum
\ No newline at end of file
diff --git a/docs/source/guides/Getting_Started.ipynb b/docs/source/guides/Getting_Started.ipynb
index dc65b34..8a77ef2 100644
--- a/docs/source/guides/Getting_Started.ipynb
+++ b/docs/source/guides/Getting_Started.ipynb
@@ -107,7 +107,7 @@
     "def prior(amplitude, length_scale, noise):\n",
     "    return models.gaussian_process(\n",
     "        means.zero(),\n",
-    "        kernels.scale(nn.softplus(amplitude), kernels.rbf(nn.softplus(length_scale))),\n",
+    "        kernels.scaled(kernels.rbf(nn.softplus(length_scale)), nn.softplus(amplitude)),\n",
     "        nn.softplus(noise),\n",
     "    )"
    ]
@@ -115,16 +115,6 @@
   {
    "cell_type": "code",
    "execution_count": 6,
-   "id": "38ee0023-19ed-4c4b-b35f-49e3b8c3647e",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "optimizer = optax.adam(0.01)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 7,
    "id": "d664e642-509b-4000-bb46-85605f477133",
    "metadata": {},
    "outputs": [],
@@ -156,6 +146,16 @@
     "}"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "38ee0023-19ed-4c4b-b35f-49e3b8c3647e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "optimizer = optax.adam(0.01)"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 9,
@@ -197,7 +197,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "id": "c468936c-6aea-41c6-a794-e7a885da0bd4",
    "metadata": {},
    "outputs": [],
@@ -206,31 +206,31 @@
     "    x_train,\n",
     "    y_train,\n",
     "    means.zero(),\n",
-    "    kernels.scale(nn.softplus(next_params['amplitude']), kernels.rbf(nn.softplus(next_params['length_scale']))),\n",
+    "    kernels.scaled(kernels.rbf(nn.softplus(next_params['length_scale'])), nn.softplus(next_params['amplitude'])),\n",
     "    nn.softplus(next_params['noise']),\n",
     ")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 16,
+   "execution_count": 14,
    "id": "f31b1562-5cc6-4a84-9f9e-8ee0f6f0c92a",
    "metadata": {},
    "outputs": [],
    "source": [
-    "acqf = acquisitions.upper_confidence_bound(2.0, surrogate)\n",
+    "acqf = acquisitions.upper_confidence_bound(surrogate, beta=2.0)\n",
     "maximizer = maximizers.bfgs(bounds, q=1, num_restarts=25, num_raw_samples=500)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 17,
+   "execution_count": 16,
    "id": "a0e6c4c9-dd2e-481c-a446-101ead7c1cf9",
    "metadata": {},
    "outputs": [],
    "source": [
-    "init_candidates = maximizer.init(maximizer_key, acqf)\n",
-    "candidates, values = maximizer.maximize(init_candidates, acqf)"
+    "candidates = maximizer.init(maximizer_key, acqf)\n",
+    "next_candidates, values = maximizer.maximize(candidates, acqf)"
    ]
   }
  ],
diff --git a/docs/source/index.rst b/docs/source/index.rst
index ab84e40..f36dab6 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -15,6 +15,7 @@ It comes with a low-level interfaces for:
       * Surrogate Models
    * **Constructing and optimizing acquisition functions** (`boax.optimization`)
       * Acquisition Functions
+      * Constraint Functions
       * Maximizers
       * Samplers
 


Test patch:
diff --git a/tests/optimization/acquisitions_test.py b/tests/optimization/acquisitions_test.py
index 56f27f8..c355743 100644
--- a/tests/optimization/acquisitions_test.py
+++ b/tests/optimization/acquisitions_test.py
@@ -4,7 +4,7 @@
 from jax import random
 
 from boax.optimization import acquisitions
-from boax.util import const
+from boax.utils.functools import const
 
 
 class AcquisitionsTest(parameterized.TestCase):
@@ -13,13 +13,12 @@ def test_expected_improvement(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
-    best = 0.0
     model = const((loc, cov))
+    best = 0.0
     candidates = jnp.empty(())
 
-    ei = acquisitions.expected_improvement(best, model)(candidates)
-    lei = acquisitions.log_expected_improvement(best, model)(candidates)
+    ei = acquisitions.expected_improvement(model, best)(candidates)
+    lei = acquisitions.log_expected_improvement(model, best)(candidates)
 
     np.testing.assert_allclose(jnp.log(ei), lei, atol=1e-4)
 
@@ -28,13 +27,12 @@ def test_probability_of_improvement(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
-    best = 1.96
     model = const((loc, cov))
+    best = 1.96
     candidates = jnp.empty(())
 
-    pi = acquisitions.probability_of_improvement(best, model)(candidates)
-    lpi = acquisitions.log_probability_of_improvement(best, model)(candidates)
+    pi = acquisitions.probability_of_improvement(model, best)(candidates)
+    lpi = acquisitions.log_probability_of_improvement(model, best)(candidates)
 
     np.testing.assert_allclose(jnp.log(pi), lpi, atol=1e-4)
 
@@ -43,12 +41,11 @@ def test_upper_confidence_bound(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
-    beta = 1.0
     model = const((loc, cov))
+    beta = 1.0
     candidates = jnp.empty(())
 
-    ucb = acquisitions.upper_confidence_bound(beta, model)(candidates)
+    ucb = acquisitions.upper_confidence_bound(model, beta)(candidates)
     expected = loc + jnp.sqrt(jnp.diag(cov))
 
     np.testing.assert_allclose(ucb, expected, atol=1e-4)
@@ -58,7 +55,6 @@ def test_posterior(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
     model = const((loc, cov))
     candidates = jnp.empty(())
 
@@ -73,13 +69,12 @@ def test_q_expected_improvement(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
-    best = 0.0
-    base_samples = random.normal(key3, (10,))
     model = const((loc, cov))
+    base_samples = random.normal(key3, (10,))
+    best = 0.0
     candidates = jnp.empty(())
 
-    qei = acquisitions.q_expected_improvement(best, base_samples, model)(
+    qei = acquisitions.q_expected_improvement(model, base_samples, best)(
       candidates
     )
 
@@ -90,15 +85,14 @@ def test_q_probability_of_improvement(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
-    best = 0.0
-    tau = 1.0
-    base_samples = random.normal(key3, (10,))
     model = const((loc, cov))
+    base_samples = random.normal(key3, (10,))
+    tau = 1.0
+    best = 0.0
     candidates = jnp.empty(())
 
     qei = acquisitions.q_probability_of_improvement(
-      best, tau, base_samples, model
+      model, base_samples, tau, best
     )(candidates)
 
     self.assertEqual(qei.shape, (10,))
@@ -108,18 +102,45 @@ def test_q_upper_confidence_bound(self):
 
     loc = random.uniform(key1, (10,))
     cov = random.uniform(key2, (10,)) * jnp.identity(10)
-
-    beta = 2.0
-    base_samples = random.normal(key3, (10,))
     model = const((loc, cov))
+    base_samples = random.normal(key3, (10,))
+    beta = 2.0
     candidates = jnp.empty(())
 
-    qei = acquisitions.q_upper_confidence_bound(beta, base_samples, model)(
+    qei = acquisitions.q_upper_confidence_bound(model, base_samples, beta)(
       candidates
     )
 
     self.assertEqual(qei.shape, (10,))
 
+  def test_constrained(self):
+    key1, key2, key3, key4 = random.split(random.key(0), 4)
+
+    loc = random.uniform(key1, (10,))
+    cov = random.uniform(key2, (10,)) * jnp.identity(10)
+    c1 = random.uniform(key3, (10,))
+    c2 = random.uniform(key4, (10,))
+    model = const((loc, cov))
+    best = 0.0
+    candidates = jnp.empty(())
+
+    cei = acquisitions.constrained(
+      acquisitions.expected_improvement(model, best),
+      const(c1),
+      const(c2),
+    )(candidates)
+
+    clei = acquisitions.log_constrained(
+      acquisitions.log_expected_improvement(model, best),
+      const(jnp.log(c1)),
+      const(jnp.log(c2)),
+    )(candidates)
+
+    print(cei)
+    print(clei)
+
+    np.testing.assert_allclose(jnp.log(cei), clei, atol=1e-4)
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/optimization/constraints_test.py b/tests/optimization/constraints_test.py
new file mode 100644
index 0000000..24e7841
--- /dev/null
+++ b/tests/optimization/constraints_test.py
@@ -0,0 +1,32 @@
+import numpy as np
+from absl.testing import absltest, parameterized
+from jax import numpy as jnp
+from jax import random
+
+from boax.optimization import constraints
+from boax.utils.functools import const
+
+
+class ConstraintsTest(parameterized.TestCase):
+  def test_range(self):
+    key1, key2 = random.split(random.key(0))
+
+    loc = random.uniform(key1, (10,))
+    cov = random.uniform(key2, (10,)) * jnp.identity(10)
+    model = const((loc, cov))
+    lower = 0.2
+    upper = 0.8
+    candidates = jnp.empty(())
+
+    le = constraints.less_or_equal(model, lower)(candidates)
+    lle = constraints.log_less_or_equal(model, lower)(candidates)
+
+    ge = constraints.greater_or_equal(model, upper)(candidates)
+    lge = constraints.log_greater_or_equal(model, upper)(candidates)
+
+    np.testing.assert_allclose(jnp.log(le), lle, atol=1e-4)
+    np.testing.assert_allclose(jnp.log(ge), lge, atol=1e-4)
+
+
+if __name__ == '__main__':
+  absltest.main()
diff --git a/tests/optimization/maximizers_test.py b/tests/optimization/maximizers_test.py
index ce44224..594d201 100644
--- a/tests/optimization/maximizers_test.py
+++ b/tests/optimization/maximizers_test.py
@@ -3,7 +3,7 @@
 from jax import random
 
 from boax.optimization import maximizers
-from boax.util import const
+from boax.utils.functools import const
 
 
 class MaximizersTest(parameterized.TestCase):
diff --git a/tests/prediction/kernels_test.py b/tests/prediction/kernels_test.py
index d8e4770..9e68d54 100644
--- a/tests/prediction/kernels_test.py
+++ b/tests/prediction/kernels_test.py
@@ -1,5 +1,3 @@
-from functools import partial
-
 import numpy as np
 from absl.testing import absltest, parameterized
 from jax import numpy as jnp
@@ -127,9 +125,9 @@ def test_periodic_kernel(self):
 
     self.assertEqual(result.shape, (10, 10))
 
-  def test_scale(self):
+  def test_scaled(self):
     key1, key2 = random.split(random.key(0))
-    amplitude = 5
+    amplitude = 5.0
     length_scale = 0.2
 
     x = random.uniform(key1, shape=(10,))
@@ -137,22 +135,34 @@ def test_scale(self):
 
     inner = kernels.rbf(length_scale)
 
-    result = kernels.scale(amplitude, inner)(x, y)
+    result = kernels.scaled(inner, amplitude)(x, y)
     expected = amplitude * inner(x, y)
 
     np.testing.assert_allclose(result, expected, atol=1e-4)
 
-  def test_combine(self):
+  def test_sum(self):
     key1, key2 = random.split(random.key(0))
-    length_scales = [0.2, 0.3, 0.4]
 
     x = random.uniform(key1, shape=(10,))
     y = random.uniform(key2, shape=(10,))
 
-    inner = list(map(kernels.rbf, length_scales))
+    result = kernels.sum(*map(kernels.rbf, [0.2, 0.3, 0.4]))(x, y)
+    expected = (
+      kernels.rbf(0.2)(x, y) + kernels.rbf(0.3)(x, y) + kernels.rbf(0.4)(x, y)
+    )
+
+    np.testing.assert_allclose(result, expected, atol=1e-4)
+
+  def test_product(self):
+    key1, key2 = random.split(random.key(0))
+
+    x = random.uniform(key1, shape=(10,))
+    y = random.uniform(key2, shape=(10,))
 
-    result = kernels.combine(partial(jnp.sum, axis=0), *inner)(x, y)
-    expected = sum(kernel(x, y) for kernel in inner)
+    result = kernels.product(*map(kernels.rbf, [0.2, 0.3, 0.4]))(x, y)
+    expected = (
+      kernels.rbf(0.2)(x, y) * kernels.rbf(0.3)(x, y) * kernels.rbf(0.4)(x, y)
+    )
 
     np.testing.assert_allclose(result, expected, atol=1e-4)