chunking_enabled

chunking_enabled(enabled=True, chunk_size=None)

Context manager to temporarily enable/disable chunked FFT conv.

Example:

from nvsubquadratic.ops.fftconv_chunked import (
    fftconv2d_fp32_bhl,
    chunking_enabled,
)

# Use chunking with chunk_size=32
with chunking_enabled(True, chunk_size=32):
    y = fftconv2d_fp32_bhl(x, kernel)

# Disable chunking (use standard impl)
with chunking_enabled(False):
    y = fftconv2d_fp32_bhl(x, kernel)

Parameters:

• enabled (bool) – Whether to enable chunking within this context.

• chunk_size (int | None) – Optional chunk size override for this context.

Yields:

None