pytc.tc

JAX implementation of Transcorrelated method.

Functions

_array_nbytes	Return the byte size of an array-like object without copying.
_cache_d_on_device	Whether a persistent device cache should keep arr resident.
_cache_tc_kernels_on_device	Whether persistent TC kernels should be cached on one device.
_choose_tc_kernel_strategy	Decide how TC kernels (K1 + K3) should be kept for tile consumption.
_compute_2b_shard	Compute K terms for one device shard.
_evict_isdf_device_cache_for_id	Drop all cache entries for an ISDFTC instance (called from GC).
_get_local_device_free_bytes	Return free bytes for one local device when available.
_normalize_panel_layout	Normalize the internal tiled-axis layout selector.
_pad_axis	Pad one axis of an array with zeros up to target.
_panel_blk_overrides	Read optional process-wide panel/block-size overrides from the env.
_transpose_panel_layout	Return the layout required for the transpose partner tile.
invalidate_isdf_device_cache	Manually evict ISDFTC device-cache entries.
trim_panel	Strip JIT padding from the two variable axes of a tile.