mostlyai-engine 1.0.0__py3-none-any.whl

mostlyai/engine/_tabular/common.py

@@ -0,0 +1,29 @@
+# Copyright 2025 MOSTLY AI
+#
+# 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.
+
+import logging
+import time
+import torch
+from pathlib import Path
+
+_LOG = logging.getLogger(__name__)
+
+
+def load_model_weights(model: torch.nn.Module, path: Path, device: torch.device):
+    try:
+        t00 = time.time()
+        model.load_state_dict(torch.load(f=path, map_location=device, weights_only=True))
+        _LOG.info(f"loaded model weights in {time.time() - t00:.2f}s")
+    except Exception as e:
+        _LOG.warning(f"failed to load model weights: {e}")

mostlyai/engine/_tabular/encoding.py

@@ -0,0 +1,406 @@
+# Copyright 2025 MOSTLY AI
+#
+# 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.
+
+import logging
+import time
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+from joblib import Parallel, delayed, parallel_config, cpu_count
+
+from mostlyai.engine._common import (
+    ARGN_COLUMN,
+    ARGN_PROCESSOR,
+    ARGN_TABLE,
+    SIDX_SUB_COLUMN_PREFIX,
+    SLEN_SUB_COLUMN_PREFIX,
+    TGT,
+    encode_slen_sidx_sdec,
+    get_argn_name,
+    get_sequence_length_stats,
+    is_a_list,
+    is_sequential,
+    SDEC_SUB_COLUMN_PREFIX,
+    ProgressCallback,
+    ProgressCallbackWrapper,
+)
+from mostlyai.engine._encoding_types.tabular.categorical import encode_categorical
+from mostlyai.engine._encoding_types.tabular.character import encode_character
+from mostlyai.engine._encoding_types.tabular.datetime import encode_datetime
+from mostlyai.engine._encoding_types.tabular.itt import encode_itt
+from mostlyai.engine._encoding_types.tabular.lat_long import encode_latlong
+from mostlyai.engine._encoding_types.tabular.numeric import encode_numeric
+from mostlyai.engine.domain import ModelEncodingType
+from mostlyai.engine._workspace import Workspace, ensure_workspace_dir, reset_dir
+
+_LOG = logging.getLogger(__name__)
+
+
+def encode(
+    workspace_dir: str | Path | None = None,
+    update_progress: ProgressCallback | None = None,
+) -> None:
+    _LOG.info("ENCODE_TABULAR started")
+    t0 = time.time()
+    with ProgressCallbackWrapper(update_progress) as progress:
+        # build paths based on workspace dir
+        workspace_dir = ensure_workspace_dir(workspace_dir)
+        workspace = Workspace(workspace_dir)
+        reset_dir(workspace.encoded_data_val.path)
+        reset_dir(workspace.encoded_data_trn.path)
+        reset_dir(workspace.encoded_data_path)
+
+        has_context = workspace.ctx_data_path.exists()
+
+        tgt_pqt_partitions = workspace.tgt_data.fetch_all()
+        if has_context:
+            ctx_pqt_partitions = workspace.ctx_data.fetch_all()
+            if len(tgt_pqt_partitions) != len(ctx_pqt_partitions):
+                raise RuntimeError("partition files for tgt and ctx do not match")
+        else:
+            ctx_pqt_partitions = []
+        tgt_stats = workspace.tgt_stats.read()
+        ctx_stats = workspace.ctx_stats.read()
+
+        for i in range(len(tgt_pqt_partitions)):
+            _encode_partition(
+                tgt_partition_file=tgt_pqt_partitions[i],
+                tgt_stats=tgt_stats,
+                output_path=workspace.encoded_data_path,
+                ctx_partition_file=ctx_pqt_partitions[i] if has_context else None,
+                ctx_stats=ctx_stats if has_context else None,
+                n_jobs=min(cpu_count() - 1, 16),
+            )
+            progress.update(completed=i, total=len(tgt_pqt_partitions) + 1)
+    _LOG.info(f"ENCODE_TABULAR finished in {time.time() - t0:.2f}s")
+
+
+def _encode_partition(
+    *,
+    tgt_partition_file: Path,
+    tgt_stats: dict,
+    output_path: Path,
+    ctx_partition_file: Path | None = None,
+    ctx_stats: dict | None = None,
+    n_jobs: int = 1,
+) -> None:
+    seq_len_stats = get_sequence_length_stats(tgt_stats)
+    is_sequential = tgt_stats["is_sequential"]
+
+    tgt_context_key = tgt_stats.get("keys", {}).get("context_key")
+    ctx_primary_key = ctx_stats.get("keys", {}).get("primary_key") if ctx_stats else None
+
+    # encode target data
+    df = pd.read_parquet(tgt_partition_file)
+    df, _, tgt_context_key = encode_df(
+        df,
+        tgt_stats,
+        ctx_primary_key=None,
+        tgt_context_key=tgt_context_key,
+        n_jobs=n_jobs,
+    )
+
+    has_context = ctx_partition_file is not None and tgt_context_key and ctx_primary_key
+    ctx_stats = ctx_stats or {}
+    if has_context:
+        # check consistency of partitioned file names
+        if ctx_partition_file and tgt_partition_file.name != ctx_partition_file.name:
+            raise RuntimeError(f"mismatch: {tgt_partition_file}!={ctx_partition_file}")
+        # encode context data
+        df_ctx = pd.read_parquet(ctx_partition_file)
+        df_ctx, ctx_primary_key, _ = encode_df(
+            df_ctx,
+            ctx_stats,
+            ctx_primary_key=ctx_primary_key,
+            tgt_context_key=None,
+            n_jobs=n_jobs,
+        )
+        # pad each list with one extra item
+        df_ctx = pad_horizontally(df_ctx, padding_value=0, right=False)
+
+    if is_sequential:
+        assert isinstance(tgt_context_key, str)
+        # trim sequences to (privacy-protected) max_len
+        max_len = seq_len_stats["max"]
+        df = df[df.groupby(tgt_context_key).cumcount() < max_len].reset_index(drop=True)
+        # enrich with sequence lengths and sequence indexes
+        df = _enrich_slen_sidx_sdec(df, tgt_context_key, max_len)
+        # flatten to list columns
+        df = flatten_frame(df, tgt_context_key)
+        # add empty records for IDs, that are present in context, but not in target; i.e., for zero-sequence records
+        if has_context:
+            zero_seq_ids = list(set(df_ctx[ctx_primary_key]) - set(df[tgt_context_key]))
+            df_miss = pd.DataFrame({tgt_context_key: zero_seq_ids})
+            df_pads = pd.DataFrame({c: [[]] for c in df.columns if c != tgt_context_key})
+            df_miss = df_miss.merge(df_pads, how="cross")
+            df = pd.concat([df, df_miss], axis=0).reset_index(drop=True)
+        # pad each list with one extra item
+        df = pad_horizontally(df, padding_value=0, right=True)
+    elif has_context:
+        # add 0-rows for IDs, that are present in context, but not in target; i.e., for zero-sequence records
+        zero_seq_ids = list(set(df_ctx[ctx_primary_key]) - set(df[tgt_context_key]))
+        df_miss = pd.DataFrame({tgt_context_key: zero_seq_ids})
+        df_pads = pd.DataFrame({c: [0] for c in df.columns if c != tgt_context_key})
+        df_miss = df_miss.merge(df_pads, how="cross")
+        df = pd.concat([df, df_miss], axis=0).reset_index(drop=True)
+        # ensure that max 1 item is retained per context_id for flat mode
+        df = df[df.groupby(tgt_context_key).cumcount() < 1]
+
+    # merge context with target
+    if has_context:
+        df = df_ctx.merge(df, left_on=ctx_primary_key, right_on=tgt_context_key, how="inner")
+
+    # drop all key columns
+    keys = [tgt_context_key, ctx_primary_key]
+    df.drop(columns=keys, inplace=True, errors="ignore")
+
+    # return, if nothing to write or if no target columns are present
+    tgt_columns = [c for c in df.columns if c.startswith(TGT)]
+    if df.empty or len(tgt_columns) == 0:
+        return
+
+    # shuffle and persist to disk as parquet files
+    df = df.sample(frac=1)
+    df.to_parquet(output_path / tgt_partition_file.name, engine="pyarrow", index=False)
+    _LOG.info(f"encoded partition {tgt_partition_file.name} {df.shape}")
+
+
+def encode_df(
+    df: pd.DataFrame,
+    stats: dict,
+    ctx_primary_key: str | None = None,
+    tgt_context_key: str | None = None,
+    n_jobs: int = 1,
+) -> tuple[pd.DataFrame, str | None, str | None]:
+    """
+    Encodes a given table represented by a DataFrame object. The result will be delivered
+    as DataFrame, as well.
+
+    :param df: input data to be encoded
+    :param stats: stats for each of the columns
+    :param ctx_primary_key: context primary key
+    :param tgt_context_key: target context key
+    :return: encoded data and keys following columns' naming conventions
+    """
+
+    if ctx_primary_key and ctx_primary_key not in df:
+        raise ValueError(f"primary key `{ctx_primary_key}` not present")
+    if tgt_context_key and tgt_context_key not in df:
+        raise ValueError(f"context key `{tgt_context_key}` not present")
+    if ctx_primary_key and df[[ctx_primary_key]].duplicated().any():
+        raise ValueError("Duplicate primary keys in table")
+
+    df_columns = []
+
+    # key columns
+    context_keys = df[tgt_context_key] if tgt_context_key in df.columns else None
+    if ctx_primary_key is not None:
+        unique_primary_key = f"flt/primary_key.{ctx_primary_key}"
+        df_columns.append(df[ctx_primary_key].reset_index(drop=True).to_frame(unique_primary_key))
+        ctx_primary_key = unique_primary_key
+    if tgt_context_key is not None:
+        unique_context_key = f"context_key.{tgt_context_key}"
+        df_columns.append(df[tgt_context_key].reset_index(drop=True).to_frame(unique_context_key))
+        tgt_context_key = unique_context_key
+
+    # data columns
+    stats = stats or {"columns": {}}
+    delayed_encodes = []
+    for column in stats["columns"].keys():
+        if column not in df:
+            # skip encoding for columns that are not present in the data
+            continue
+        column_stats = stats["columns"][column]
+        if not set(column_stats.keys()) - {"encoding_type"}:
+            # all partitions are empty; skip encoding
+            continue
+        # encode (empty and non-empty) column
+        delayed_encodes.append(
+            delayed(_encode_col)(
+                values=df[column],
+                column_stats=column_stats,
+                context_keys=context_keys,
+            )
+        )
+    if delayed_encodes:
+        with parallel_config("loky", n_jobs=n_jobs):
+            df_columns.extend(Parallel()(delayed_encodes))
+
+    df = pd.concat(df_columns, axis=1) if df_columns else pd.DataFrame()
+    return df, ctx_primary_key, tgt_context_key
+
+
+def _encode_col(
+    values: pd.Series,
+    column_stats: dict,
+    context_keys: pd.Series | None = None,
+) -> pd.DataFrame:
+    is_sequential_column = is_sequential(values)
+    if is_sequential_column:
+        # explode nested columns and encode the same way as flat columns
+        non_empties = values.apply(lambda v: len(v) if is_a_list(v) else 1) > 0
+        # generate serial context_keys, if context_keys are not provided
+        context_keys = (
+            context_keys
+            if context_keys is not None
+            else pd.Series(range(len(values)), index=values.index).rename("__ckey")
+        )
+        # explode non-empty values and context_keys in sync, remember sequence keys, reset index afterwards
+        df = pd.concat([values[non_empties], context_keys[non_empties]], axis=1)
+        df = df.explode(values.name)
+        # trim sequences to max_seq_len
+        max_seq_len = column_stats.get("seq_len", {}).get("max")
+        if max_seq_len:
+            df = df[df.groupby(context_keys.name).cumcount() < max_seq_len]
+        sequence_keys = df.index
+        df = df.reset_index(drop=True)
+        values, context_keys = df[values.name], df[context_keys.name]
+
+    df = _encode_flat_col(
+        encoding_type=column_stats["encoding_type"],
+        values=values,
+        column_stats=column_stats,
+        context_keys=context_keys,
+    )
+
+    if is_sequential_column:
+        # flatten encoded, non-empty sequences
+        encoded_sequences = flatten_frame(df.assign(__sequence_key=sequence_keys), "__sequence_key").set_index(
+            "__sequence_key"
+        )
+        # inject empty sequences back into their original positions
+        empty_sequences = pd.DataFrame(index=non_empties.index[~non_empties], columns=encoded_sequences.columns).map(
+            lambda _: []
+        )
+        df = pd.concat([encoded_sequences, empty_sequences]).sort_index()
+
+    df.reset_index(drop=True, inplace=True)
+
+    if column_stats.get(ARGN_TABLE) is not None:
+        # assign sanitized, unique column names for usage in ARGN
+        df.columns = [
+            get_argn_name(
+                argn_processor=column_stats[ARGN_PROCESSOR],
+                argn_table=column_stats[ARGN_TABLE],
+                argn_column=column_stats[ARGN_COLUMN],
+                argn_sub_column=sub_column,
+            )
+            for sub_column in df.columns
+        ]
+
+    return df
+
+
+def _encode_flat_col(
+    encoding_type: ModelEncodingType,
+    values: pd.Series,
+    column_stats: dict,
+    context_keys: pd.Series | None,
+) -> pd.DataFrame:
+    if encoding_type == ModelEncodingType.tabular_categorical:
+        df = encode_categorical(values, column_stats, context_keys)
+    elif encoding_type in [
+        ModelEncodingType.tabular_numeric_discrete,
+        ModelEncodingType.tabular_numeric_binned,
+        ModelEncodingType.tabular_numeric_digit,
+    ]:
+        df = encode_numeric(values, column_stats, context_keys)
+    elif encoding_type == ModelEncodingType.tabular_datetime:
+        df = encode_datetime(values, column_stats, context_keys)
+    elif encoding_type == ModelEncodingType.tabular_datetime_relative:
+        df = encode_itt(values, column_stats, context_keys)
+    elif encoding_type == ModelEncodingType.tabular_character:
+        df = encode_character(values, column_stats, context_keys)
+    elif encoding_type == ModelEncodingType.tabular_lat_long:
+        df = encode_latlong(values, column_stats, context_keys)
+    else:
+        raise RuntimeError(f"unknown encoding_type `{encoding_type}`")
+    return df
+
+
+def flatten_frame(df: pd.DataFrame, group_key: str) -> pd.DataFrame:
+    """
+    Flattens non-key columns. See also `explode_frame` for reverse.
+
+    This is an optimized flattening of data to rows by a given foreign key column. It is equivalent to
+    `return data.groupby(foreign_key_column).agg(list).reset_index(level=0)`, but was benchmarked being
+     significantly faster
+
+    :example:
+        this method converts
+           key  product  is_paid
+             1        3        0
+             1        2        1
+             2        9        1
+        to
+           key  product  is_paid
+             1   [3, 2]   [0, 1]
+             2      [9]      [1]
+    """
+    orig_columns_order = df.columns
+    sorted_data = (
+        df.reset_index(drop=True).rename_axis("__flatten_index").sort_values(by=[group_key, "__flatten_index"])
+    )  # pre-group rows by FK, and preserve the order of their appearance by the index
+
+    keys = sorted_data.pop(group_key).values.T
+    values = sorted_data.values.T
+    column_names = sorted_data.columns.values
+    unique_keys, index = np.unique(keys, True)  # cutoff indexes for the groups (by FK)
+    grouped_values = [
+        np.split(value, index[1:]) for value in values
+    ]  # a list of lists of arrays: column -> rows as a list -> grouped array of values
+
+    flattened_data = {group_key: unique_keys}
+    flattened_data.update(
+        {
+            column_name: [list(group) for group in column_values]
+            for column_name, column_values in zip(column_names, grouped_values)
+        }
+    )
+
+    flattened_data = (
+        pd.DataFrame(flattened_data).reindex(columns=orig_columns_order)
+        if len(unique_keys) > 0
+        else pd.DataFrame(columns=orig_columns_order)
+    )
+    return flattened_data
+
+
+def _enrich_slen_sidx_sdec(df: pd.DataFrame, context_key: str, max_seq_len: int) -> pd.DataFrame:
+    df = df.reset_index(drop=True)
+    slen = df.groupby(context_key)[context_key].transform("size")  # sequence length
+    sidx = df.groupby(context_key).cumcount(ascending=True)  # sequence index
+    sdec = (10 * sidx / slen.clip(lower=1)).astype(int)  # sequence index decile
+    slen = encode_slen_sidx_sdec(slen, max_seq_len=max_seq_len, prefix=SLEN_SUB_COLUMN_PREFIX)
+    sidx = encode_slen_sidx_sdec(sidx, max_seq_len=max_seq_len, prefix=SIDX_SUB_COLUMN_PREFIX)
+    sdec = encode_slen_sidx_sdec(sdec, max_seq_len=max_seq_len, prefix=SDEC_SUB_COLUMN_PREFIX)
+    df = pd.concat([slen, sidx, sdec, df], axis=1)
+    return df
+
+
+def pad_horizontally(df: pd.DataFrame, padding_value: int, right=True) -> pd.DataFrame:
+    if df.shape[0] == 0:
+        return df
+    list_cols = [c for c in df.columns if is_a_list(df.loc[0, c])]
+
+    def pad_right(x):
+        return x + [padding_value] if len(x) == 0 else x
+
+    def pad_left(x):
+        return [padding_value] + x if len(x) == 0 else x
+
+    for col in list_cols:
+        df[col] = df[col].apply(pad_right if right else pad_left)
+    return df

mostlyai/engine/_tabular/fairness.py

@@ -0,0 +1,245 @@
+# Copyright 2025 MOSTLY AI
+#
+# 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.
+
+import logging
+from functools import partial
+from typing import Any, TypedDict
+from collections.abc import Callable
+
+import pandas as pd
+import numpy as np
+
+from itertools import product
+import torch
+
+from mostlyai.engine._common import (
+    get_argn_name,
+    ARGN_COLUMN,
+    ARGN_PROCESSOR,
+    ARGN_TABLE,
+)
+from mostlyai.engine._encoding_types.tabular.categorical import CATEGORICAL_SUB_COL_SUFFIX, CATEGORICAL_UNKNOWN_TOKEN
+from mostlyai.engine.domain import ModelEncodingType, FairnessConfig
+
+_LOG = logging.getLogger(__name__)
+
+
+class FairnessTransforms(TypedDict):
+    target_sub_col: str
+    sensitive_sub_cols: list[str]
+    transforms: dict[int, dict[tuple[int], callable]]
+
+
+def _get_sensitive_groups(
+    column_stats: dict[str, Any],
+    sensitive_cols: list[str],
+    sensitive_sub_cols: list[str],
+) -> pd.DataFrame:
+    """
+    Create a DataFrame with all possible combinations of (encoded) values for the sensitive columns.
+    """
+    category_map = {}
+    for sensitive_col in sensitive_cols:
+        category_map[sensitive_col] = list(column_stats[sensitive_col]["codes"].values())
+        # when there are no rare categories, we can exclude _RARE_ to reduce the number of combinations
+        if column_stats[sensitive_col]["no_of_rare_categories"] == 0:
+            category_map[sensitive_col].remove(column_stats[sensitive_col]["codes"][CATEGORICAL_UNKNOWN_TOKEN])
+    groups_df = pd.DataFrame(list(product(*category_map.values())), columns=sensitive_sub_cols)
+    _LOG.info(f"fairness: {groups_df.shape[0]} sensitive groups")
+    return groups_df
+
+
+def _estimate_target_distribution(
+    forward_fn: Callable,
+    device: str | torch.device,
+    target_sub_col: str,
+    sample_size: int,
+    n_quantiles: int = 101,
+    seed_values: dict[str, int] | None = None,
+) -> np.ndarray:
+    """
+    Estimate the (conditional or unconditional) empirical  distribution of the target column by fetching model probabilities for
+     that column, and then computing the quantiles of that empirical distribution.
+
+    :return: A numpy array of shape (n_quantiles, target_cardinality).
+    """
+    if seed_values is None:
+        fixed_values = None
+        _LOG.info(f"fairness: sample {sample_size:,} unconditionally")
+    else:
+        fixed_values = {
+            col: torch.as_tensor([values] * sample_size, device=device).type(torch.int)
+            for col, values in seed_values.items()
+        }
+        _LOG.info(f"fairness: sample {sample_size:,} for sensitive group {seed_values}")
+
+    # fetch probs for target column
+    _, probs = forward_fn(
+        x={},
+        mode="gen",
+        batch_size=sample_size,
+        fixed_values=fixed_values,
+        return_probs=[target_sub_col],
+    )
+    # convert to numpy
+    probs = probs[target_sub_col].detach().cpu().numpy()
+    # calculate quantiles, and wrap with 0 a
+    quantiles = np.concatenate(
+        [
+            np.zeros((1, probs.shape[1])),
+            np.quantile(probs, np.linspace(0, 1, n_quantiles), axis=0),
+            np.ones((1, probs.shape[1])),
+        ]
+    )
+    return quantiles
+
+
+def get_fairness_transforms(
+    forward_fn: Callable,
+    device: str | torch.device,
+    fairness: FairnessConfig,
+    tgt_stats: dict[str, Any],
+) -> FairnessTransforms:
+    """
+    Generate the transforms for the target column conditioned on the sensitive columns.
+    """
+    _LOG.info(
+        f"fairness: create transforms for target `{fairness.target_column}` across "
+        f"sensitive columns: {''.join(fairness.sensitive_columns)}"
+    )
+
+    # validate the fairness configuration
+    if fairness.sensitive_columns is None or len(fairness.sensitive_columns) == 0:
+        raise ValueError("The sensitive columns must be specified in the fairness configuration")
+    if fairness.target_column is None:
+        raise ValueError("The target column must be specified in the fairness configuration")
+    if fairness.target_column in fairness.sensitive_columns:
+        raise ValueError("The target column cannot be a sensitive column")
+    cols = [fairness.target_column] + fairness.sensitive_columns
+    for col in cols:
+        if col not in tgt_stats["columns"]:
+            raise ValueError(f"Invalid fairness column: `{col}` not found in the training data")
+        if tgt_stats["columns"][col]["encoding_type"] != ModelEncodingType.tabular_categorical:
+            raise ValueError(f"Invalid fairness column: `{col}` is not encoded categorically")
+
+    # estimate the marginal target distribution
+    target_sub_col = get_argn_name(
+        argn_processor=tgt_stats["columns"][fairness.target_column][ARGN_PROCESSOR],
+        argn_table=tgt_stats["columns"][fairness.target_column][ARGN_TABLE],
+        argn_column=tgt_stats["columns"][fairness.target_column][ARGN_COLUMN],
+        argn_sub_column=CATEGORICAL_SUB_COL_SUFFIX,
+    )
+    marginal_target_quantiles: np.ndarray = _estimate_target_distribution(
+        forward_fn=forward_fn,
+        device=device,
+        target_sub_col=target_sub_col,
+        sample_size=15_000,
+        n_quantiles=101,
+        seed_values=None,
+    )
+
+    # estimate the conditional target distribution for each combination of sensitive values
+    sensitive_sub_cols = [
+        get_argn_name(
+            argn_processor=tgt_stats["columns"][col][ARGN_PROCESSOR],
+            argn_table=tgt_stats["columns"][col][ARGN_TABLE],
+            argn_column=tgt_stats["columns"][col][ARGN_COLUMN],
+            argn_sub_column=CATEGORICAL_SUB_COL_SUFFIX,
+        )
+        for col in fairness.sensitive_columns
+    ]
+    sensitive_groups_df = _get_sensitive_groups(
+        column_stats=tgt_stats["columns"],
+        sensitive_cols=fairness.sensitive_columns,
+        sensitive_sub_cols=sensitive_sub_cols,
+    )
+    conditional_target_quantiles: dict[tuple[int], np.ndarray] = {}
+    for _, sensitive_group in sensitive_groups_df.iterrows():
+        seed_values = sensitive_group.to_dict()
+        conditional_target_quantiles[tuple(seed_values.values())] = _estimate_target_distribution(
+            forward_fn=forward_fn,
+            device=device,
+            target_sub_col=target_sub_col,
+            sample_size=5_000,
+            n_quantiles=101,
+            seed_values=seed_values,
+        )
+
+    # calculate the transforms, one for each target value and sensitive group
+    transforms: dict[int, dict[tuple[int], callable]] = {}
+    for i in range(marginal_target_quantiles.shape[1]):
+        transforms[i] = {}
+        for sensitive_group, target_quantiles in conditional_target_quantiles.items():
+            transforms[i][sensitive_group] = partial(
+                torch_interp,
+                xp=torch.as_tensor(target_quantiles[:, i], device=device),
+                fp=torch.as_tensor(marginal_target_quantiles[:, i], device=device),
+            )
+
+    _LOG.info("fairness: created transforms")
+    return {
+        "target_sub_col": target_sub_col,
+        "sensitive_sub_cols": sensitive_sub_cols,
+        "transforms": transforms,
+    }
+
+
+def apply_fairness_transforms(
+    sub_col: str, probs: torch.Tensor, outputs: dict[str, torch.Tensor], fairness_transforms: FairnessTransforms
+) -> torch.Tensor:
+    """
+    Apply fairness transforms to the generated probabilities of the target column.
+
+    :param sub_col: The current sub column.
+    :param probs: The probabilities for the current sub column.
+    :param outputs: The sampled values for the sensitive columns.
+    :param fairness_transforms: The transform callables to apply.
+    :return: The transformed probabilities for the target column.
+    """
+    # only transform the target column
+    device = probs.device
+    if sub_col == fairness_transforms["target_sub_col"]:
+        sensitive_values = torch.stack([outputs[col] for col in fairness_transforms["sensitive_sub_cols"]]).t()
+        for i, i_transforms in fairness_transforms["transforms"].items():
+            for group, transform in i_transforms.items():
+                probs[:, i] = torch.where(
+                    torch.all(
+                        sensitive_values == torch.as_tensor(group, device=device).tile((probs.shape[0], 1)),
+                        dim=1,
+                    ),
+                    transform(probs[:, i]),
+                    probs[:, i],
+                )
+        # normalize transformed probabilities to make sure they sum to 1
+        probs /= probs.sum(dim=1, keepdim=True)
+    return probs
+
+
+def torch_interp(x: torch.Tensor, xp: torch.Tensor, fp: torch.Tensor) -> torch.Tensor:
+    """
+    Torch implementation of numpy.interp.
+    """
+    # find the indices of the bins to which each value in x belongs
+    indices = torch.searchsorted(xp, x, side="right")
+    indices = torch.clamp(indices, 1, len(xp) - 1)
+    # compute the slopes of the segments
+    x_lo = xp[indices - 1]
+    x_hi = xp[indices]
+    y_lo = fp[indices - 1]
+    y_hi = fp[indices]
+    # compute the interpolated values
+    eps = 1e-10
+    slopes = (y_hi - y_lo) / (x_hi - x_lo + eps)
+    interpolated_values = y_lo + slopes * (x - x_lo)
+    return interpolated_values