Parallel Scaling Law

from typing import List, Dict

# Scaling law with shared exponents across groups and group-specific coefficients.
# Formula: loss = L_g + A_g * num_params**(-alpha) + B_g * parallel_size**(-beta)
# Fitted values (from analysis on /app/data):
#   alpha = 0.25, beta = 0.5
#   pile: L=1.3602632039069371, A=98.21459269821332, B=0.101969095947293
#   stack: L=0.7368400879171292, A=55.921667192849526, B=0.06640252659247498

ALPHA = 0.25
BETA = 0.5

GROUP_PARAMS = {
    "pile": {
        "L": 1.3602632039069371,
        "A": 98.21459269821332,
        "B": 0.101969095947293,
    },
    "stack": {
        "L": 0.7368400879171292,
        "A": 55.921667192849526,
        "B": 0.06640252659247498,
    },
}

# Fallback parameters (average of known groups) in case an unknown group is requested
if GROUP_PARAMS:
    _Ls = [v["L"] for v in GROUP_PARAMS.values()]
    _As = [v["A"] for v in GROUP_PARAMS.values()]
    _Bs = [v["B"] for v in GROUP_PARAMS.values()]
    DEFAULT_PARAMS = {"L": sum(_Ls) / len(_Ls), "A": sum(_As) / len(_As), "B": sum(_Bs) / len(_Bs)}
else:
    DEFAULT_PARAMS = {"L": 0.0, "A": 0.0, "B": 0.0}


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law must be the same for all groups,
                but the constant parameters/coefficients can differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s).
    """
    params = GROUP_PARAMS.get(group, DEFAULT_PARAMS)
    Lg, Ag, Bg = params["L"], params["A"], params["B"]

    outputs: List[Dict[str, float]] = []
    for row in input_data:
        n = float(row.get("num_params", 0.0))
        p = float(row.get("parallel_size", 1.0))
        # Avoid invalid values
        n = max(n, 1.0)
        p = max(p, 1.0)
        pred_loss = Lg + Ag * (n ** (-ALPHA)) + Bg * (p ** (-BETA))
        outputs.append({"loss": float(pred_loss)})
    return outputs

from __future__ import annotations

import math
from typing import Dict, List

# Per-group coefficients for the log-linear scaling law:
# loss = a_g + b_g * ln(num_params) + c_g * ln(parallel_size)
# Coefficients estimated by ordinary least squares on the provided dataset.
_COEFS: Dict[str, tuple[float, float, float]] = {
    # group: (a_g, b_g, c_g)
    "pile": (4.618399622902, -0.125421944189, -0.036842959428),
    "stack": (2.599633430732, -0.071389434896, -0.023902322958),
}

# Fallback (average of known groups) if an unseen group name is provided.
_FALLBACK: tuple[float, float, float] = (
    (4.618399622902 + 2.599633430732) / 2.0,
    (-0.125421944189 + -0.071389434896) / 2.0,
    (-0.036842959428 + -0.023902322958) / 2.0,
)


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Functional form (shared across groups):
        loss = a_g + b_g * ln(num_params) + c_g * ln(parallel_size)

    Args:
        input_data: List of dicts, each containing at least 'num_params' and 'parallel_size'.
        group: Experimental group name to select the group's coefficients.

    Returns:
        A list of dictionaries, each with a single key 'loss'.
    """
    a, b, c = _COEFS.get(group, _FALLBACK)

    out: List[Dict[str, float]] = []
    for row in input_data:
        if 'num_params' not in row or 'parallel_size' not in row:
            raise KeyError("Each input row must contain 'num_params' and 'parallel_size'.")
        N = float(row['num_params'])
        P = float(row['parallel_size'])
        if N <= 0 or P <= 0:
            raise ValueError('num_params and parallel_size must be positive.')
        y = a + b * math.log(N) + c * math.log(P)
        out.append({'loss': float(y)})
    return out

from __future__ import annotations

from typing import List, Dict

# Discovered scaling law (shared exponents across groups, group-specific amplitude):
# loss = A_g * num_params**b * parallel_size**c
# Fitted on /app/data via linear regression on log(loss) with shared slopes and
# group-specific intercepts.

# Shared exponents
_B = -0.06544312287561933  # exponent for num_params
_C = -0.02049343910349895  # exponent for parallel_size

# Group-specific amplitudes A_g
_AMPLITUDE = {
    "pile": 7.844746369924655,
    "stack": 4.339456524197749,
}

# Fallback amplitude for unknown groups: geometric mean of known amplitudes
_default_amp = (1.0)
if _AMPLITUDE:
    import math
    vals = list(_AMPLITUDE.values())
    _default_amp = math.exp(sum(math.log(v) for v in vals) / len(vals))


def _safe_positive(x: float) -> float:
    # Guard against non-positive inputs (shouldn't occur in valid data)
    return x if x > 0 else 1e-12


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    The law uses a multiplicative power-law with shared exponents across groups
    and a group-specific amplitude (A_g):
        loss = A_g * num_params**b * parallel_size**c

    Args:
        input_data: List of data points with keys 'num_params' and 'parallel_size'.
        group: Name of the experimental group (e.g., 'stack', 'pile'). Determines A_g.

    Returns:
        A list with dictionaries containing {'loss': predicted_value} for each input.
    """
    A_g = _AMPLITUDE.get(group, _default_amp)
    b, c = _B, _C

    out: List[Dict[str, float]] = []
    for row in input_data:
        N = float(row.get("num_params", 0.0))
        P = float(row.get("parallel_size", 0.0))
        N = _safe_positive(N)
        P = _safe_positive(P)
        loss = A_g * (N ** b) * (P ** c)
        out.append({"loss": float(loss)})
    return out

from __future__ import annotations

from typing import Dict, List
import math

# Pre-fitted parameters from pooled log-log regression on the provided dataset.
# Functional form (shared across groups):
#   loss = A_g * (num_params)^{-alpha} * (parallel_size)^{-beta}
# Group-specific amplitude A_g; shared exponents alpha, beta.

_ALPHA = 0.06544312287562275
_BETA = 0.020493439103498882
_A_BY_GROUP: Dict[str, float] = {
    "stack": 4.339456524197764,
    "pile": 7.844746369924683,
}
# Fallback amplitude for unknown groups (geometric mean of known A's)
_A_DEFAULT = math.sqrt(_A_BY_GROUP["stack"] * _A_BY_GROUP["pile"])


def _predict_one(num_params: float, parallel_size: float, A: float) -> float:
    # Guard against non-positive inputs
    n = max(float(num_params), 1.0)
    p = max(float(parallel_size), 1.0)
    return A * (n ** (-_ALPHA)) * (p ** (-_BETA))


def law(input_data: List[Dict[str, float]], group: str) -> List[Dict[str, float]]:
    """
    Predicts output variables based on input variables according to a discovered scaling law.

    Args:
        input_data: A list of dictionaries, where each dictionary is a single data
                    point containing input variable names as keys and their
                    corresponding values. Required keys: 'num_params', 'parallel_size'.
        group: The name of the experimental group for which to make predictions.
                The functional form of the law is the same for all groups, while
                the amplitude parameter A may differ per group.

    Returns:
        A list of dictionaries, corresponding to the input_data list, with each
        dictionary containing the predicted output variable(s) under key 'loss'.
    """
    A = _A_BY_GROUP.get(group, _A_DEFAULT)
    outputs: List[Dict[str, float]] = []
    for row in input_data:
        n = row.get("num_params", 0.0)
        p = row.get("parallel_size", 1.0)
        y = _predict_one(n, p, A)
        outputs.append({"loss": float(y)})
    return outputs