dynamax

Dynamax uses jaxtyping for type declarations of JAX arrays. While these declarations are available in the documentation, run-time type checking is currently disabled.

To interpret the API signatures, use the following shorthand terminology:

Basic Types

• Scalar: An int, float, or a JAX array with shape ().
• PyTree: Any valid JAX PyTree.
• Array: A JAX array.

Array Shapes and Dtypes

Array annotations specify dimensions and data types using the following syntax:

• Named Dimensions: Array["dim1", "dim2"] represents a JAX array with shape (dim1, dim2). When dimensions share the same name across different arrays in an API call, they must have matching sizes.
• Zero-dimensional Arrays: Array[()] represents an array with shape ().
• Ellipsis: ... represents an arbitrary number of dimensions (e.g., Array["times", ...]).
• Dtypes: Array[bool] refers to a JAX array with a Boolean dtype. Other dtypes follow the same pattern.
• Combined Annotations: You can combine shape and dtype, such as Array["dim1", "dim2", bool].
• Unions: Union types are used when an argument can take different shapes (e.g., a transition matrix that can be either a constant matrix or a time-varying sequence).

To create a custom emission model in Dynamax, you must implement a class that inherits from HMMEmissions. Your class must implement the following interface:

• Properties:
  • emission_shape: Returns the shape of the emission distribution.
  • inputs_shape: Returns the shape of the inputs to the emission distribution.
• Methods:
  • initialize(key, method, ...): Returns a tuple of (params, props). params contains the actual parameter values, while props contains ParameterProperties (e.g., using TFP bijectors to enforce constraints like non-negativity during optimization).
  • log_prior(params): Computes the log prior probability of the parameters. If not implemented, the base class assumes a zero log prior.
  • distribution(params, state, inputs): Returns a tensorflow_probability.substrates.jax.distributions object representing the likelihood $p(y_t \mid x_t, z_t)$ for a given state and input.

Example of a Poisson GLM emission implementation:

class PoissonGLMHMMEmissions(HMMEmissions):
    def __init__(self, num_states, emission_dim, input_dim, ...):
        super().__init__(...)
        self.num_states = num_states
        self.emission_dim = emission_dim
        self.input_dim = input_dim

    @property
    def emission_shape(self) -> Tuple:
        return (self.emission_dim,)

    @property
    def inputs_shape(self) -> Tuple[int]:
        return (self.input_dim,)

    def initialize(self, key, method="prior", ...):
        # ... implementation returning (params, props) ...
        # Use tfb.Softplus() in props to enforce non-negativity
        props = ParamsPoissonGLMHMMEmissions(weights=ParameterProperties(constrainer=tfb.Softplus()))
        return params, props

    def distribution(self, params, state, inputs):
        activations = params.weights[state] @ inputs
        return tfd.Independent(tfd.Poisson(rate=activations), 1)

To create a new state space model in Dynamax, you must subclass the SSM abstract base class and implement the following required methods and properties:

Required Abstract Methods

• initial_distribution(params, inputs): Returns a tfd.Distribution representing the initial state $p(z_1 \mid \theta)$.
• transition_distribution(params, state, inputs): Returns a tfd.Distribution representing the conditional next state $p(z_{t+1} \mid z_t, u_t, \theta)$.
• emission_distribution(params, state, inputs): Returns a tfd.Distribution representing the conditional emission $p(y_t \mid z_t, u_t, \theta)$.

Required Properties

• emission_shape: A property returning a tuple (or pytree of tuples) specifying the shape of a single time step's emissions.
• inputs_shape (Optional): A property returning the shape of a single time step's inputs. Defaults to None if there are no inputs.

Optional Methods

• log_prior(params): Returns the log prior probability of the parameters. Defaults to 0.0 if not implemented.

from dynamax.ssm import SSM

class MyCustomSSM(SSM):
    @property
    def emission_shape(self):
        return (D,)

    def initial_distribution(self, params, inputs):
        # return a tensorflow_probability distribution
        pass

    def transition_distribution(self, params, state, inputs):
        # return a tensorflow_probability distribution
        pass

    def emission_distribution(self, params, state, inputs=None):
        # return a tensorflow_probability distribution
        pass

Once you have defined your custom HMMEmissions class, you can assemble a full HMM by inheriting from the base HMM class. This allows you to combine your custom emissions with standard components like StandardHMMInitialState and StandardHMMTransitions.

In your custom HMM subclass:

1. __init__: Initialize the component objects (initial state, transitions, and your custom emissions) and call super().__init__(num_states, initial_component, transition_component, emission_component).
2. inputs_shape property: Define the shape of the inputs expected by your emission model.
3. initialize method: Delegate parameter initialization to the underlying components and return a combined HMMParameterSet and HMMPropertySet.

By following this pattern, your custom model automatically inherits powerful methods like .sample(), .fit_em(), and .smoother() from the base HMM class.

class PoissonGLMHMM(HMM):
    def __init__(self, num_states, emission_dim, input_dim, ...):
        self.inputs_dim = input_dim
        initial_component = StandardHMMInitialState(num_states, ...)
        transition_component = StandardHMMTransitions(num_states, ...)
        emission_component = PoissonGLMHMMEmissions(num_states, emission_dim, input_dim, ...)
        super().__init__(num_states, initial_component, transition_component, emission_component)

    @property
    def inputs_shape(self) -> Tuple[int, ...]:
        return (self.inputs_dim,)

    def initialize(self, key, method="prior", ...):
        # ... split key and call component.initialize() ...
        return ParamsPoissonGLMHMM(**params), ParamsPoissonGLMHMM(**props)