from collections.abc import Sequence
from typing import TYPE_CHECKING, Literal
import jax
import jax.flatten_util
import jax.numpy as jnp
import numpy as np
from jax.typing import ArrayLike
from ..types import PositionInput, PyTree
from .model import Model
if TYPE_CHECKING:
from collections.abc import Callable
[docs]
class LogProb:
"""
Interface for evaluating the unnormalized log probability of a Liesel model.
Also provides access to the first and second derivatives.
Position mappings are not modified. Derivative methods convert the mapping to a
dictionary before differentiation. When applying JAX transformations externally,
pass a dictionary or another registered pytree to the transformed function.
Parameters
----------
model
A Liesel model instance.
component
Which component of the model's log probability to evaluate.
diff_mode
Which auto-diff mode to use for the Hessian.
See Also
---------
~liesel.model.FlatLogProb : A similar class that returns gradients and hessians as
arrays.
Examples
--------
We initialize a very basic Liesel model:
>>> x = lsl.Var.new_param(
... jnp.zeros(2),
... dist=lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
... name="x",
... )
>>> model = lsl.Model([x])
Now we initialize the log prob object:
>>> lp = lsl.LogProb(model)
And evaluate the log prob (the unnormalized log posterior) at a new position:
>>> lp({"x": jnp.array([1.0, 2.0])})
Array(-4.3378773, dtype=float32)
Now we evaluate the gradient of the unnormalized log posterior at the new position:
>>> lp.grad({"x": jnp.array([1.0, 2.0])})
{'x': Array([-1., -2.], dtype=float32)}
And, finally, the hessian:
>>> lp.hessian({"x": jnp.array([1.0, 2.0])})
{'x': {'x': Array([[-1., -0.],
[-0., -1.]], dtype=float32)}}
"""
if TYPE_CHECKING:
component: Literal["log_prob", "log_lik", "log_prior"]
"""Which component of the model's log probability to evaluate."""
diff_mode: Literal["forward", "reverse"]
"""Which auto-diff mode to use for the Hessian."""
model: Model
"""A Liesel model instance."""
def __init__(
self,
model: Model,
component: Literal["log_prob", "log_lik", "log_prior"] = "log_prob",
diff_mode: Literal["forward", "reverse"] = "forward",
):
self.model = model
self._grad_fn = jax.grad(self.log_prob)
if diff_mode == "forward":
self._hessian_fn = jax.jacfwd(self._grad_fn)
elif diff_mode == "reverse":
self._hessian_fn = jax.jacrev(self._grad_fn)
else:
raise ValueError(f"Unrecognized argument value {diff_mode=}")
self.component = component
self.diff_mode = diff_mode
def __call__(self, position: PositionInput) -> float | np.number | jax.Array:
"""
Log probability function evaluated at provided ``position``.
"""
return self.log_prob(position=position)
[docs]
def log_prob(self, position: PositionInput) -> float | np.number | jax.Array:
"""
Log probability function evaluated at provided ``position``.
"""
updated_state = self.model.update_state(position, self.model.state)
return updated_state[f"_model_{self.component}"].value
[docs]
def grad(self, position: PositionInput) -> dict[str, PyTree]:
"""
Gradient of the log probability function with respect to the ``position``.
"""
return self._grad_fn(dict(position))
[docs]
def hessian(
self,
position: PositionInput,
) -> dict[str, PyTree]:
"""
Hessian of the log probability function with respect to the ``position``.
"""
return self._hessian_fn(dict(position))
[docs]
class FlatLogProb:
"""
Interface for evaluating the unnormalized log probability of a Liesel model.
Also provides access to the first and second derivatives.
The methods :meth:`FlatLogProb.grad <liesel.model.FlatLogProb.grad>` and
:meth:`FlatLogProb.hessian <liesel.model.FlatLogProb.hessian>` are
flattened, which means the expect arrays as inputs and return arrays.
Parameters
----------
model
A Liesel model instance.
position_keys
Names of the variables at which to evaluate the log probability. Other \
variables will be kept fixed at their current values in the model state.
component
Which component of the model's log probability to evaluate.
diff_mode
Which auto-diff mode to use for the Hessian.
See Also
--------
~liesel.model.LogProb : A similar class that returns gradients and hessians as
dictionaries.
Examples
--------
We initialize a very basic Liesel model:
>>> x = lsl.Var.new_param(
... jnp.zeros(2),
... dist=lsl.Dist(tfd.Normal, loc=0.0, scale=1.0),
... name="x",
... )
>>> model = lsl.Model([x])
Now we initialize the log prob object:
>>> lp = lsl.FlatLogProb(model, ["x"])
And an array of new values to evaluate the log probability at:
>>> xnew = jnp.array([1.0, 2.0])
>>> lp(xnew)
Array(-4.3378773, dtype=float32)
>>> lp.grad(xnew)
Array([-1., -2.], dtype=float32)
>>> lp.hessian(xnew)
Array([[-1., -0.],
[-0., -1.]], dtype=float32)
"""
if TYPE_CHECKING:
component: Literal["log_prob", "log_lik", "log_prior"]
"""Which component of the model's log probability to evaluate."""
diff_mode: Literal["forward", "reverse"]
"""Which auto-diff mode to use for the Hessian."""
model: Model
"""A Liesel model instance."""
unravel_fn: Callable
"""Function mapping a flat parameter vector back to its position dictionary."""
def __init__(
self,
model: Model,
position_keys: Sequence[str],
component: Literal["log_prob", "log_lik", "log_prior"] = "log_prob",
diff_mode: Literal["forward", "reverse"] = "forward",
):
self.model = model
if position_keys is None:
position_keys = [
var.name for var in self.model.vars.values() if var.parameter
]
position = self.model.extract_position(position_keys, self.model.state)
_, unravel_fn = jax.flatten_util.ravel_pytree(position)
self.unravel_fn = unravel_fn
self._grad_fn = jax.grad(self)
if diff_mode == "forward":
self._hessian_fn = jax.jacfwd(self._grad_fn)
elif diff_mode == "reverse":
self._hessian_fn = jax.jacrev(self._grad_fn)
else:
raise ValueError(f"Unrecognized argument value {diff_mode=}")
self.component = component
self.diff_mode = diff_mode
def __call__(self, flat_position: ArrayLike) -> float | np.number | jax.Array:
"""
Log probability function evaluated at provided ``flat_position``.
"""
return self.log_prob(flat_position=flat_position)
[docs]
def log_prob(self, flat_position: ArrayLike) -> float | np.number | jax.Array:
"""
Log probability function evaluated at provided ``flat_position``.
"""
position = self.unravel_fn(jnp.asarray(flat_position))
updated_state = self.model.update_state(position, self.model.state)
return updated_state[f"_model_{self.component}"].value
[docs]
def grad(self, flat_position: ArrayLike) -> jax.Array:
"""
Gradient of the log probability function with respect to the ``flat_position``.
"""
return self._grad_fn(jnp.asarray(flat_position))
[docs]
def hessian(self, flat_position: ArrayLike) -> jax.Array:
"""
Hessian of the log probability function with respect to the ``flat_position``.
"""
return self._hessian_fn(jnp.asarray(flat_position))
