p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/train_utils.py (python)


def print_compiled_memory_stats(compiled_stats):
    """from: https://github.com/AI-Hypercomputer/maxtext/blob/b18829fbaa48aec7ac350a03e62248e24c6a76b2/MaxText/max_utils.py#L739"""
    output_gb = bytes_to_gb(compiled_stats.output_size_in_bytes)
    temp_gb = bytes_to_gb(compiled_stats.temp_size_in_bytes)
    argument_gb = bytes_to_gb(compiled_stats.argument_size_in_bytes)
    alias_gb = bytes_to_gb(compiled_stats.alias_size_in_bytes)
    host_temp_gb = bytes_to_gb(compiled_stats.host_temp_size_in_bytes)
    total_gb = output_gb + temp_gb + argument_gb - alias_gb
    print(
        f"Total memory size: {total_gb:.1f} GB, Output size: {output_gb:.1f} GB, Temp size: {temp_gb:.1f} GB, "
        f"Argument size: {argument_gb:.1f} GB, Host temp size: {host_temp_gb:.1f} GB."
    )


def print_compiled_cost_analysis(cost_stats):
    flops = float(cost_stats.get("flops", 0.0))
    bytes_accessed = float(cost_stats.get("bytes accessed", 0.0))
    gb = bytes_to_gb(bytes_accessed) if bytes_accessed else 0.0
    intensity = (flops / bytes_accessed) if bytes_accessed else float("nan")
    print(
        f"FLOPs: {flops:.3e}, Bytes: {bytes_accessed:.3e} ({gb:.1f} GB), "
        f"Intensity: {intensity:.1f} FLOPs/byte"
    )