MoE Scaling Law

import numpy as np

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).
    """
    predictions = []
    # In this specific dataset, there's only one group: 'all_data'.
    # If there were multiple groups, we would load different parameters based on the 'group' argument.
    if group == 'all_data':
        # Fitted parameters from the analysis
        A = -0.0274019225275534
        B = -0.0705124860796179
        C = 10.069179203081296
    else:
        # Handle cases where the group is not recognized or parameters are not available
        # For this problem, 'all_data' is the only expected group.
        # A more robust solution might raise an error or return NaNs/default values.
        raise ValueError(f"Group '{group}' not recognized or parameters not available.")

    for data_point in input_data:
        num_experts = data_point['num_experts']
        dense_parameter_count = data_point['dense_parameter_count']

        # Ensure num_experts and dense_parameter_count are positive for the power law
        if num_experts <= 0 or dense_parameter_count <= 0:
            # Handle invalid input, e.g., by returning NaN or a default value
            predicted_loss = np.nan
        else:
            predicted_loss = C * (num_experts**A) * (dense_parameter_count**B)
        
        predictions.append({'loss_validation': predicted_loss})
    
    return predictions

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).
    """
    # Fitted parameters for the 'all_data' group
    # A_coeff: -0.02740192252755339
    # B_coeff: -0.07051248607961787
    # C_coeff: 10.069179203081283
    
    # In a more complex scenario with multiple groups, you would store these parameters
    # in a dictionary keyed by group name and retrieve them here.
    # Since there's only one group ('all_data'), we can hardcode for simplicity, 
    # but keep the group argument for future extensibility.

    C = 10.069179203081283
    A = -0.02740192252755339
    B = -0.07051248607961787

    predictions = []
    for data_point in input_data:
        num_experts = data_point['num_experts']
        dense_parameter_count = data_point['dense_parameter_count']

        # Apply the scaling law
        predicted_loss_validation = C * (num_experts ** A) * (dense_parameter_count ** B)
        predictions.append({'loss_validation': predicted_loss_validation})

    return predictions

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).
    """
    # Fitted parameters for the 'all_data' group
    # C: 10.324241375460156
    # a (exponent for num_experts): -0.0279704190795864
    # b (exponent for dense_parameter_count): -0.07178446289739228

    # Since the problem statement implies there could be multiple groups, 
    # but our dataset only contained 'all_data', we'll hardcode parameters 
    # for 'all_data'. If more groups were present, this would be a dictionary lookup.
    if group == 'all_data':
        C = 10.324241375460156
        a = -0.0279704190795864
        b = -0.07178446289739228
    else:
        raise ValueError(f"Unknown group: {group}")

    predictions = []
    for data_point in input_data:
        num_experts = data_point['num_experts']
        dense_parameter_count = data_point['dense_parameter_count']

        # Apply the scaling law
        predicted_loss_validation = C * (num_experts**a) * (dense_parameter_count**b)
        predictions.append({'loss_validation': predicted_loss_validation})

    return predictions

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).
    """
    import numpy as np

    # Fitted parameters for the 'all_data' group
    # These values were obtained by fitting the scaling law to the provided dataset.
    # Mathematical form: loss_validation = a + b * log(num_experts) + c * log(dense_parameter_count)
    fitted_params = {
        'all_data': {
            'a': 6.077175446130047,
            'b': -0.07132315536715568,
            'c': -0.18021292228947294
        }
    }

    if group not in fitted_params:
        raise ValueError(f"No fitted parameters found for group: {group}")

    a = fitted_params[group]['a']
    b = fitted_params[group]['b']
    c = fitted_params[group]['c']

    predictions = []
    for data_point in input_data:
        num_experts = data_point['num_experts']
        dense_parameter_count = data_point['dense_parameter_count']

        # Ensure non-zero values for log transformation
        num_experts_log = np.log(num_experts) if num_experts > 0 else np.log(1e-9)
        dense_parameter_count_log = np.log(dense_parameter_count) if dense_parameter_count > 0 else np.log(1e-9)
        
        predicted_loss_validation = a + b * num_experts_log + c * dense_parameter_count_log
        predictions.append({'loss_validation': predicted_loss_validation})

    return predictions