Factorial Types

def __call__(self, model_inputs: ArrayTreeLike) -> ArrayLike:
    """Extract the factorial indices from model inputs.

    Args:
        model_inputs: Model inputs.

    Returns:
        Indices of the factors to extract. Integer array.
    """
    ...

def __call__(
    self,
    factorial_state: ArrayTreeLike,
    factorial_inds: ArrayLike,
) -> ArrayTree:
    """Extract factors from factorial state.

    E.g. factorial_state might encode factorial `means` with shape (F, d) and
    `chol_covs` with shape (F, d, d). Then `model_inputs` tells us factors `i` and
    `j` are relevant, so we extract `means[i]` and `means[j]` and `chol_covs[i]` and
    `chol_covs[j]`. Thus we return `means` with shape (2, d) and `chol_covs` with
    shape (2, d, d).

    factorial_state: Factorial Kalman state storing means and chol_covs
        with shape (F, d) and (F, d, d) respectively.
    factorial_inds: Indices of the factors to extract. Integer array.
        factorial_inds.ndim == 0 removes the factorial dimension and extracts
            a single factor.
        factorial_inds.ndim == 1 retains the factorial dimension,
            even if len(factorial_inds) == 1.

    Returns:
        Local factorial state with factorial dimension of length len(factorial_inds).
            If factorial_inds is a single integer, the local factorial state should
            not have a factorial dimension.
    """
    ...

def __call__(
    self,
    local_factorial_state: ArrayTreeLike,
) -> ArrayTree:
    """Extract factors from factorial state and combine into a joint local state.

    E.g., local_factorial_state might encode factorial `means` with shape (2, d)
    and `chol_covs` with shape (2, d, d),
    which is then combined into a joint state with shape (2 * d,)
    and block diagonal `joint_chol_cov` with shape (2 * d, 2 * d).

    Args:
        local_factorial_state: Factorial state with factorial index as the first
            dimension. Typically contains only a small number of factors, as it's
            applied after an `Extract` operation.

    Returns:
        Joint state with no factorial index dimension.
    """
    ...

def __call__(
    self,
    local_state: ArrayTree,
    num_factors: int,
) -> ArrayTree:
    """Marginalize joint state into factored state.

    E.g. `local_state` might have shape (2 * d,) and `joint_chol_cov`
    with shape (2 * d, 2 * d). Then we marginalize out the joint local state into
    two factorial `means` with shape (2, d) and `chol_covs` with shape (2, d, d).

    Args:
        local_state: Joint local state with no factorial index dimension.
        num_factors: Number of factors to marginalize out. Integer.
            This is typically equal to len(factorial_inds).

    Returns:
        Factorial state with factorial index as the first dimension and
            `num_factors` factors (length of first dimension).
    """
    ...

def __call__(
    self,
    local_factorial_state: ArrayTree,
    factorial_state: ArrayTree,
    factorial_inds: ArrayLike,
) -> ArrayTree:
    """Marginalize joint state into factored state and insert into factorial state.

    E.g. `local_factorial_state` might have shape (2, d) and `joint_chol_cov`
    with shape (2, d, d). Then we insert `means[0]` and `means[1]` into
    `state[i]` and `state[j]` respectively. Similarly, we insert `chol_covs[0]` and
    `chol_covs[1]`. In both cases, we overwrite the existing factors in the
    factorial state for `i` and `j`, leaving the other factors unchanged.
    Here `i` and `j` are determined from `factorial_inds`.

    Args:
        local_factorial_state: Local factorial state with factorial index as the first
            dimension and `len(factorial_inds)` factors (length of first dimension).
        factorial_state: Factorial state with factorial index as the first dimension.
        factorial_inds: Indices of the factors to insert. Integer array.
            factorial_inds.ndim == 0 will be treated the same as
            factorial_inds.ndim == 1 with len(factorial_inds) == 1
            (i.e. insert a single factor).

    Returns:
        Factorial state with factorial index as the first dimension.
            The updated factors are inserted into the factorial state.
            The remaining factors are left unchanged.
    """
    ...

def __call__(
    self, init_state: ArrayTreeLike, model_inputs: ArrayTreeLike
) -> ArrayTree:
    """Any processing of the output of `init_prepare` for factorial inference.

    Args:
        init_state: Output from base inference method's `init_prepare`
        model_inputs: The model inputs at the first time point.
    """
    ...

def extract_and_join(
    self, factorial_state: ArrayTreeLike, model_inputs: ArrayTreeLike
) -> ArrayTree:
    """Extract and join the relevant factors into a joint local state.

    Args:
        factorial_state: Factorial state with factorial index as the first dimension.
        model_inputs: Model inputs, from which the factorial indices are extracted.

    Returns:
        Joint local state with no factorial index dimension.
    """
    factorial_inds = self.get_factorial_indices(model_inputs)
    factorial_inds = jnp.asarray(factorial_inds)

    assert factorial_inds.ndim == 1, (
        "factorial_inds must be a 1D array to be used with join"
    )

    local_factorial_state = self.extract(factorial_state, factorial_inds)
    return self.join(local_factorial_state)

def marginalize_and_insert(
    self,
    local_state: ArrayTree,
    factorial_state: ArrayTree,
    model_inputs: ArrayTreeLike,
) -> ArrayTree:
    """Marginalize and insert the relevant factors into a factorial state.

    Args:
        local_state: Joint local state with no factorial index dimension.
        factorial_state: Factorial state with factorial index as the first dimension.
        model_inputs: Model inputs, from which the factorial indices are extracted.

    Returns:
        Factorial state with factorial index as the first dimension.
            The updated factors are inserted into the factorial state.
            The remaining factors are left unchanged.
    """
    factorial_inds = self.get_factorial_indices(model_inputs)
    factorial_inds = jnp.asarray(factorial_inds)
    num_factors = len(factorial_inds)
    local_factorial_state = self.marginalize(local_state, num_factors)
    return self.insert(local_factorial_state, factorial_state, factorial_inds)