Performance¤

Module: generative_models.core.performance

Source: generative_models/core/performance.py

Overview¤

Performance infrastructure for roofline analysis and optimization.

This module provides core performance analysis capabilities including:

All implementations follow JAX/Flax NNX best practices and avoid numpy usage within any performance-critical code paths.

class HardwareDetector

class HardwareSpecs

class PerformanceEstimator

class RooflineMetrics

def __init__()

def __init__()

def analyze_roofline()

def benchmark_operation()

def calculate_arithmetic_intensity()

def detect_hardware()

def estimate_flops_attention()

def estimate_flops_linear()

def estimate_memory_usage()

def estimate_transformer_layer_performance()

def get_batch_size_recommendation()

def get_critical_batch_size()

def get_optimal_batch_size()

def is_batch_size_optimal()

def profile_jax_function()