nlsq.memory_pool module

Memory pool for optimization algorithms.

This module provides memory pool allocation to reduce overhead from repeated array allocations in optimization loops.

Key Features (Task Group 5): - Size-class bucketing: Round shapes to nearest 1KB/10KB/100KB for 5x reuse increase - Reuse statistics tracking: Monitor reuse_rate = reused_allocations / total_allocations - Adaptive sizing: Small arrays (1KB buckets), medium (10KB), large (100KB)

nlsq.caching.memory_pool.round_to_bucket(nbytes)[source]

Round memory size to nearest bucket for better pool reuse.

Uses tiered bucketing strategy (Task 5.4): - Small arrays (<10KB): Round to nearest 1KB - Medium arrays (10KB-100KB): Round to nearest 10KB - Large arrays (>100KB): Round to nearest 100KB

Parameters:

nbytes (int) – Memory size in bytes

Returns:

bucketed_bytes – Rounded memory size for bucketing

Return type:

int

Examples

>>> round_to_bucket(800)    # Small array
1024                         # Rounded to 1KB
>>> round_to_bucket(8500)   # Medium array
10240                        # Rounded to 10KB
>>> round_to_bucket(85000)  # Large array
102400                       # Rounded to 100KB
class nlsq.caching.memory_pool.MemoryPool(max_pool_size=10, enable_stats=False, enable_bucketing=True)[source]

Bases: object

Memory pool for reusable array buffers.

Pre-allocates buffers for common array shapes to avoid repeated allocations during optimization iterations.

pools

Dictionary mapping (shape, dtype) to list of available buffers

Type:

dict

allocated

Dictionary tracking allocated buffers

Type:

dict

max_pool_size

Maximum number of buffers per shape/dtype combination

Type:

int

stats

Statistics on pool usage

Type:

dict

__init__(max_pool_size=10, enable_stats=False, enable_bucketing=True)[source]

Initialize memory pool.

Parameters:

  • max_pool_size (int) – Maximum number of buffers to keep per shape/dtype

  • enable_stats (bool) – Track allocation statistics

  • enable_bucketing (bool) – Enable size-class bucketing for better reuse (Task 5.4)

allocate(shape, dtype=<class 'jax.numpy.float64'>)[source]

Allocate array from pool or create new one.

Parameters:

  • shape (tuple) – Shape of array to allocate

  • dtype (type) – Data type of array

Returns:

array – Allocated array (may be reused from pool)

Return type:

jnp.ndarray

Notes

When bucketing is enabled, arrays are pooled by size classes (1KB/10KB/100KB) for better reuse rates (Task 5.4).

release(arr)[source]

Return array to pool for reuse.

Parameters:

arr (jnp.ndarray) – Array to return to pool

Notes

When bucketing is enabled, arrays are stored in size-class buckets for better reuse (Task 5.4).

clear()[source]

Clear all pools and reset statistics.

get_stats()[source]

Get pool usage statistics.

Returns:

stats – Pool usage statistics including reuse_rate (Task 5.5)

Return type:

dict

Notes

reuse_rate = reused_allocations / total_allocations With bucketing enabled, expect 5x higher reuse rates.

__enter__()[source]

Context manager entry.

__exit__(exc_type, exc_val, exc_tb)[source]

Context manager exit - clear pool.

class nlsq.caching.memory_pool.TRFMemoryPool(m, n, dtype=<class 'jax.numpy.float64'>)[source]

Bases: object

Specialized memory pool for Trust Region Reflective algorithm.

Pre-allocates buffers for common TRF operations.

Parameters:

  • m (int) – Number of residuals

  • n (int) – Number of parameters

  • dtype (type) – Data type for arrays

__init__(m, n, dtype=<class 'jax.numpy.float64'>)[source]

Initialize TRF memory pool.

Parameters:

  • m (int) – Number of residuals

  • n (int) – Number of parameters

  • dtype (type) – Data type

get_jacobian_buffer()[source]

Get Jacobian buffer (m×n).

get_residual_buffer()[source]

Get residual buffer (m).

get_gradient_buffer()[source]

Get gradient buffer (n).

get_step_buffer()[source]

Get step buffer (n).

get_x_buffer()[source]

Get parameter buffer (n).

reset()[source]

Reset all buffers to zero.

nlsq.caching.memory_pool.get_global_pool(enable_stats=False)[source]

Get or create global memory pool.

Parameters:

enable_stats (bool) – Enable statistics tracking

Returns:

pool – Global memory pool instance

Return type:

MemoryPool

nlsq.caching.memory_pool.clear_global_pool()[source]

Clear the global memory pool.

Notes

For test isolation, this resets the global pool to None, forcing fresh initialization on next access.

Overview

The memory_pool module provides memory pool allocation to reduce overhead from repeated array allocations in optimization loops.

Key Features (Task Group 5)

  • Size-class bucketing: Round shapes to nearest 1KB/10KB/100KB for 5x reuse increase

  • Reuse statistics tracking: Monitor reuse_rate = reused_allocations / total_allocations

  • Adaptive sizing: Small arrays (1KB buckets), medium (10KB), large (100KB)

  • Memory efficiency: Pre-allocate buffers for common array shapes

Performance Targets

  • Peak memory reduction: 12.5% reduction in memory usage

  • Reuse rate: 90% reuse rate via size-class bucketing

  • Allocation overhead: Minimize malloc/free calls in hot paths

Classes

class nlsq.caching.memory_pool.MemoryPool(max_pool_size=10, enable_stats=False, enable_bucketing=True)[source]

Bases: object

Memory pool for reusable array buffers.

Pre-allocates buffers for common array shapes to avoid repeated allocations during optimization iterations.

pools

Dictionary mapping (shape, dtype) to list of available buffers

Type:

dict

allocated

Dictionary tracking allocated buffers

Type:

dict

max_pool_size

Maximum number of buffers per shape/dtype combination

Type:

int

stats

Statistics on pool usage

Type:

dict

__init__(max_pool_size=10, enable_stats=False, enable_bucketing=True)[source]

Initialize memory pool.

Parameters:

  • max_pool_size (int) – Maximum number of buffers to keep per shape/dtype

  • enable_stats (bool) – Track allocation statistics

  • enable_bucketing (bool) – Enable size-class bucketing for better reuse (Task 5.4)

allocate(shape, dtype=<class 'jax.numpy.float64'>)[source]

Allocate array from pool or create new one.

Parameters:

  • shape (tuple) – Shape of array to allocate

  • dtype (type) – Data type of array

Returns:

array – Allocated array (may be reused from pool)

Return type:

jnp.ndarray

Notes

When bucketing is enabled, arrays are pooled by size classes (1KB/10KB/100KB) for better reuse rates (Task 5.4).

release(arr)[source]

Return array to pool for reuse.

Parameters:

arr (jnp.ndarray) – Array to return to pool

Notes

When bucketing is enabled, arrays are stored in size-class buckets for better reuse (Task 5.4).

clear()[source]

Clear all pools and reset statistics.

get_stats()[source]

Get pool usage statistics.

Returns:

stats – Pool usage statistics including reuse_rate (Task 5.5)

Return type:

dict

Notes

reuse_rate = reused_allocations / total_allocations With bucketing enabled, expect 5x higher reuse rates.

__enter__()[source]

Context manager entry.

__exit__(exc_type, exc_val, exc_tb)[source]

Context manager exit - clear pool.

Functions

nlsq.caching.memory_pool.round_to_bucket(nbytes)[source]

Round memory size to nearest bucket for better pool reuse.

Uses tiered bucketing strategy (Task 5.4): - Small arrays (<10KB): Round to nearest 1KB - Medium arrays (10KB-100KB): Round to nearest 10KB - Large arrays (>100KB): Round to nearest 100KB

Parameters:

nbytes (int) – Memory size in bytes

Returns:

bucketed_bytes – Rounded memory size for bucketing

Return type:

int

Examples

>>> round_to_bucket(800)    # Small array
1024                         # Rounded to 1KB
>>> round_to_bucket(8500)   # Medium array
10240                        # Rounded to 10KB
>>> round_to_bucket(85000)  # Large array
102400                       # Rounded to 100KB

Example Usage

from nlsq.caching.memory_pool import MemoryPool
import jax.numpy as jnp

# Create memory pool with statistics tracking
pool = MemoryPool(max_pool_size=10, enable_stats=True, enable_bucketing=True)

# Allocate array from pool
shape = (100, 10)
dtype = jnp.float32
array = pool.allocate(shape, dtype)

# Use array in computation
result = array * 2.0

# Return to pool for reuse
pool.release(array)

# Check pool statistics
stats = pool.get_stats()
print(f"Reuse rate: {stats['reuse_rate']:.2%}")
print(f"Total allocations: {stats['total_allocations']}")

Bucketing Strategy

The memory pool uses tiered bucketing for better reuse:

  • Small arrays (<10KB): Round to nearest 1KB

  • Medium arrays (10KB-100KB): Round to nearest 10KB

  • Large arrays (>100KB): Round to nearest 100KB

from nlsq.caching.memory_pool import round_to_bucket

print(round_to_bucket(800))  # Small: 1024 (1KB)
print(round_to_bucket(8500))  # Medium: 10240 (10KB)
print(round_to_bucket(85000))  # Large: 102400 (100KB)

See Also