create_dim_handler

pymc_extras.prior.create_dim_handler(desired_dims: Dims) → DimHandler

Wrap the handle_dims() function to always use the same desired_dims.

Parameters:

desired_dims (Dims) – The desired dimensions to align to.

Returns:

A function that takes a tensor and its current dims and aligns it to the desired dims.

Return type:

DimHandler

Examples

Create a dim handler to align to (“channel”, “group”).

import numpy as np

from pymc_extras.prior import create_dim_handler

dim_handler = create_dim_handler(("channel", "group"))

result = dim_handler(np.array([1, 2, 3]), dims="channel")