transport.samples.kernels.cg_python_batched
module torchmatch.transport.samples.kernels.cg_python_batchedPython-level batched CG for HVP acceleration.
This module implements a batched CG solver that uses the
apply_plan_vec_shifted kernels for improved performance and reduced memory
usage.
This avoids the mysterious Triton compilation state issues observed with
the fully-inline batched CG kernel while still providing a clean interface
for batched CG solving.
Includes a torch.compile-compatible version for additional speedup.
Members
type
- PythonBatchedCgInfoInformation about Python batched CG execution.
function
- compiled_cg_solveSolve H @ x = b using CG with torch.compile optimization.
- python_batched_cg_solveSolve H @ x = b using CG with external apply_plan_vec kernels.
Source: torchmatch/transport/samples/kernels/cg_python_batched.py:1