pytc.tc

JAX implementation of Transcorrelated method.

Functions

_cache_d_on_device	Whether a persistent device cache should keep arr resident.
_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).
_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.
trim_panel	Strip JIT padding from the two variable axes of a tile.