Module `gabenet.utils`

Expand source code

from copy import deepcopy
import re

from scipy.stats import binom  # type: ignore


def freeze_trainable_states(train_states) -> tuple[dict, dict]:
    """Lift out "training-mode" states to `params`.

    Args:
        train_states: States (pytree) of a model during training
            mode (i.e., `is_training=True`).

    Returns:
        A haiku-compatible params, states pair with the frozen states moved to params.
    """
    states = deepcopy(train_states)
    params: dict = {}
    for layer_name, layer_params in states.items():
        if layer_name not in params:
            params[layer_name] = {}
        if "poisson_layer" in layer_name:
            params[layer_name]["phi"] = layer_params.pop("phi")
        elif "cap_layer" in layer_name:
            params[layer_name]["r"] = layer_params.pop("r")
        elif "gamma_layer" in layer_name:
            params[layer_name]["phi"] = layer_params.pop("phi")
        else:
            raise KeyError(f"Unknown layer {layer_name}.")
    return params, states


def is_uniform(histogram, n_replicates, alpha: float = 0.1) -> bool:
    """Determine if histogram generated from `n_replicates` is uniformly distributed.

    Args:
        alpha: Probability of false positive.
    """
    # Since the histogram is supposed to be uniform, the count in each bucket is
    # binomially distributed.
    n_bins = len(histogram)
    kwargs = {"n": n_replicates, "p": 1 / n_bins}

    # We divide by the number of bins because one bin (by accident) may be out of the
    # bands.
    q_value = alpha / n_bins
    y_lower = binom.ppf(q=q_value / 2, **kwargs)
    y_upper = binom.ppf(q=1.0 - q_value / 2, **kwargs)

    # Validate that the counts in the binds all fall within the 1-alpha % expected
    # variation around uniform distribution.
    in_bands = (histogram < y_upper) & (histogram > y_lower)
    return all(in_bands)


def to_snake_string(camel_case: str) -> str:
    """Convert a camel case string to snake case."""
    return re.sub(r"(?<!^)(?=[A-Z])", "_", camel_case).lower()

Functions

def freeze_trainable_states(train_states) ‑> tuple[dict, dict]

Lift out "training-mode" states to params.

Args

train_states: States (pytree) of a model during training mode (i.e., is_training=True).

Returns

A haiku-compatible params, states pair with the frozen states moved to params.

Expand source code

def freeze_trainable_states(train_states) -> tuple[dict, dict]:
    """Lift out "training-mode" states to `params`.

    Args:
        train_states: States (pytree) of a model during training
            mode (i.e., `is_training=True`).

    Returns:
        A haiku-compatible params, states pair with the frozen states moved to params.
    """
    states = deepcopy(train_states)
    params: dict = {}
    for layer_name, layer_params in states.items():
        if layer_name not in params:
            params[layer_name] = {}
        if "poisson_layer" in layer_name:
            params[layer_name]["phi"] = layer_params.pop("phi")
        elif "cap_layer" in layer_name:
            params[layer_name]["r"] = layer_params.pop("r")
        elif "gamma_layer" in layer_name:
            params[layer_name]["phi"] = layer_params.pop("phi")
        else:
            raise KeyError(f"Unknown layer {layer_name}.")
    return params, states

def is_uniform(histogram, n_replicates, alpha: float = 0.1) ‑> bool

Determine if histogram generated from n_replicates is uniformly distributed.

Args

alpha: Probability of false positive.

Expand source code

def is_uniform(histogram, n_replicates, alpha: float = 0.1) -> bool:
    """Determine if histogram generated from `n_replicates` is uniformly distributed.

    Args:
        alpha: Probability of false positive.
    """
    # Since the histogram is supposed to be uniform, the count in each bucket is
    # binomially distributed.
    n_bins = len(histogram)
    kwargs = {"n": n_replicates, "p": 1 / n_bins}

    # We divide by the number of bins because one bin (by accident) may be out of the
    # bands.
    q_value = alpha / n_bins
    y_lower = binom.ppf(q=q_value / 2, **kwargs)
    y_upper = binom.ppf(q=1.0 - q_value / 2, **kwargs)

    # Validate that the counts in the binds all fall within the 1-alpha % expected
    # variation around uniform distribution.
    in_bands = (histogram < y_upper) & (histogram > y_lower)
    return all(in_bands)

def to_snake_string(camel_case: str) ‑> str

Convert a camel case string to snake case.

Expand source code

def to_snake_string(camel_case: str) -> str:
    """Convert a camel case string to snake case."""
    return re.sub(r"(?<!^)(?=[A-Z])", "_", camel_case).lower()