tomojax.torch
Optional PyTorch autograd bridge to TomoJAX's exact matched projectors.
Import this module explicitly; importing tomojax does not load PyTorch.
| Name | Summary |
|---|---|
Projector | A fixed-geometry CUDA projector and its exact transpose, with autograd. |
Projector
class Projector(geometry: ScanGeometry, *, grid: Grid | None = None)A fixed-geometry CUDA projector and its exact transpose, with autograd.
operator(volume) maps (nx, ny, nz) to (views, v, u);
operator.adjoint(projections) maps back. Inputs must be contiguous
CUDA float32 tensors with 16-byte-aligned storage (standard allocations).
Clone an unaligned offset view before passing it. The DLPack boundary
shares storage in both directions;
neither inputs nor outputs pass through the host. Noncontiguous gradients
produced by PyTorch are made contiguous before applying the transpose.
Geometry and grid are snapshotted at construction and are not trainable. Each direction compiles on first use, then reuses its compiled executable. Both directions support differentiation, including higher derivatives of compositions. Install PyTorch separately for your CUDA version; it is not required by the rest of TomoJAX.
Parameters
geometry(ScanGeometry): Fixed acquisition geometry, including detector, angles and any poses. All lengths use a common physical unit.grid(Grid or None, optional): Reconstruction grid; None uses geometry.grid. Geometry and grid are copied at construction, so later caller mutations do not affect the operator.
Projector.forward
forward(volume: torch.Tensor) -> torch.TensorProject a CUDA volume; its gradient uses the exact transpose.
Parameters
volume(torch.Tensor): Contiguous CUDA float32 volume of shape (nx, ny, nz), with 16-byte-aligned storage, on the operator's JAX-compatible device.
Returns
- torch.Tensor: CUDA float32 line integrals of shape (views, detector.nv, detector.nu).
Raises
- ValueError: If dtype, device, shape, contiguity or storage alignment is unsupported.
Projector.adjoint
adjoint(projections: torch.Tensor) -> torch.TensorBackproject with the exact transpose; its gradient uses the projector.
Parameters
projections(torch.Tensor): Contiguous CUDA float32 projections of shape (views, detector.nv, detector.nu), with 16-byte-aligned storage, on the operator's JAX-compatible device.
Returns
- torch.Tensor: CUDA float32 adjoint volume of shape (nx, ny, nz).
Raises
- ValueError: If dtype, device, shape, contiguity or storage alignment is unsupported.