TomoJAX

tomojax.torch

Optional PyTorch autograd bridge to TomoJAX's exact matched projectors.

Import this module explicitly; importing tomojax does not load PyTorch.

NameSummary
ProjectorA fixed-geometry CUDA projector and its exact transpose, with autograd.

Projector

class Projector(geometry: ScanGeometry, *, grid: Grid | None = None)
Defined in src/tomojax/torch.py:52

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.Tensor

Project 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.Tensor

Backproject 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.

On this page