diffct 2.3.0__tar.gz → 2.4.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (141) hide show
  1. {diffct-2.3.0 → diffct-2.4.0}/CHANGELOG.md +26 -0
  2. {diffct-2.3.0 → diffct-2.4.0}/CITATION.cff +5 -5
  3. {diffct-2.3.0 → diffct-2.4.0}/PKG-INFO +24 -18
  4. {diffct-2.3.0 → diffct-2.4.0}/README.md +23 -17
  5. {diffct-2.3.0 → diffct-2.4.0}/diffct/__init__.py +5 -1
  6. {diffct-2.3.0 → diffct-2.4.0}/diffct/_solver_state.py +42 -25
  7. {diffct-2.3.0 → diffct-2.4.0}/diffct/backends.py +140 -50
  8. {diffct-2.3.0 → diffct-2.4.0}/diffct/chunking.py +34 -0
  9. {diffct-2.3.0 → diffct-2.4.0}/diffct/distribution.py +67 -7
  10. {diffct-2.3.0 → diffct-2.4.0}/diffct/geometry.py +93 -17
  11. {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/cone_beam.py +21 -8
  12. diffct-2.4.0/diffct/kernels/surface_geometry.py +130 -0
  13. {diffct-2.3.0 → diffct-2.4.0}/diffct/operators.py +765 -343
  14. {diffct-2.3.0 → diffct-2.4.0}/diffct/projectors.py +61 -30
  15. {diffct-2.3.0 → diffct-2.4.0}/diffct/solvers.py +18 -11
  16. diffct-2.4.0/diffct/surface_gradients.py +117 -0
  17. diffct-2.4.0/diffct/surfaces.py +146 -0
  18. diffct-2.4.0/docs/video/README.md +244 -0
  19. {diffct-2.3.0 → diffct-2.4.0}/pyproject.toml +1 -1
  20. diffct-2.4.0/tests/test_call_trajectory.py +252 -0
  21. {diffct-2.3.0 → diffct-2.4.0}/tests/test_chunked_projector.py +40 -25
  22. diffct-2.4.0/tests/test_cone_backprojection_transpose.py +189 -0
  23. {diffct-2.3.0 → diffct-2.4.0}/tests/test_distributed_training_transforms.py +5 -0
  24. diffct-2.4.0/tests/test_geometry_conversions.py +178 -0
  25. diffct-2.4.0/tests/test_native_volume_layout.py +328 -0
  26. diffct-2.4.0/tests/test_native_volume_layout_distributed.py +79 -0
  27. diffct-2.4.0/tests/test_parallel3d.py +376 -0
  28. diffct-2.4.0/tests/test_parallel3d_surface_oracle.py +149 -0
  29. diffct-2.4.0/tests/test_plan_memory_bounds.py +252 -0
  30. diffct-2.4.0/tests/test_surface_collective_failures.py +521 -0
  31. diffct-2.4.0/tests/test_surface_composition_kernel.py +358 -0
  32. diffct-2.4.0/tests/test_surface_compute_shards.py +564 -0
  33. diffct-2.4.0/tests/test_surface_geometry_cudasim.py +514 -0
  34. diffct-2.4.0/tests/test_surface_gradient_replay.py +226 -0
  35. diffct-2.4.0/tests/test_surface_native_regressions.py +281 -0
  36. diffct-2.4.0/tests/test_surface_pipeline_contracts.py +388 -0
  37. diffct-2.4.0/tests/test_surface_trajectory_transform.py +142 -0
  38. diffct-2.4.0/tests/test_trajectory_admission_cost.py +257 -0
  39. diffct-2.4.0/tests/test_trajectory_surface_integration.py +615 -0
  40. diffct-2.3.0/diffct/surfaces.py +0 -103
  41. diffct-2.3.0/docs/video/README.md +0 -119
  42. {diffct-2.3.0 → diffct-2.4.0}/.gitignore +0 -0
  43. {diffct-2.3.0 → diffct-2.4.0}/LICENSE +0 -0
  44. {diffct-2.3.0 → diffct-2.4.0}/NOTICE +0 -0
  45. {diffct-2.3.0 → diffct-2.4.0}/diffct/_solver_tv.py +0 -0
  46. {diffct-2.3.0 → diffct-2.4.0}/diffct/analytical.py +0 -0
  47. {diffct-2.3.0 → diffct-2.4.0}/diffct/analytical_execution.py +0 -0
  48. {diffct-2.3.0 → diffct-2.4.0}/diffct/constants.py +0 -0
  49. {diffct-2.3.0 → diffct-2.4.0}/diffct/differentiable.py +0 -0
  50. {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/__init__.py +0 -0
  51. {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/analytical_transpose.py +0 -0
  52. {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/fan_beam.py +0 -0
  53. {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/geometry_vjp.py +0 -0
  54. {diffct-2.3.0 → diffct-2.4.0}/diffct/kernels/parallel_beam.py +0 -0
  55. {diffct-2.3.0 → diffct-2.4.0}/diffct/pipeline.py +0 -0
  56. {diffct-2.3.0 → diffct-2.4.0}/diffct/reconstruction.py +0 -0
  57. {diffct-2.3.0 → diffct-2.4.0}/diffct/storage.py +0 -0
  58. {diffct-2.3.0 → diffct-2.4.0}/diffct/streaming.py +0 -0
  59. {diffct-2.3.0 → diffct-2.4.0}/diffct/utils.py +0 -0
  60. {diffct-2.3.0 → diffct-2.4.0}/examples/README.md +0 -0
  61. {diffct-2.3.0 → diffct-2.4.0}/examples/_block_cgls.py +0 -0
  62. {diffct-2.3.0 → diffct-2.4.0}/examples/_common.py +0 -0
  63. {diffct-2.3.0 → diffct-2.4.0}/examples/data/NOTICE +0 -0
  64. {diffct-2.3.0 → diffct-2.4.0}/examples/data/preprocess_walnut.py +0 -0
  65. {diffct-2.3.0 → diffct-2.4.0}/examples/walnut_reconstruction.py +0 -0
  66. {diffct-2.3.0 → diffct-2.4.0}/formal/README.md +0 -0
  67. {diffct-2.3.0 → diffct-2.4.0}/pytest.ini +0 -0
  68. {diffct-2.3.0 → diffct-2.4.0}/tests/__init__.py +0 -0
  69. {diffct-2.3.0 → diffct-2.4.0}/tests/_phantoms.py +0 -0
  70. {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/README.md +0 -0
  71. {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/__init__.py +0 -0
  72. {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/_common.py +0 -0
  73. {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/test_bench_cone.py +0 -0
  74. {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/test_bench_fan.py +0 -0
  75. {diffct-2.3.0 → diffct-2.4.0}/tests/benchmarks/test_bench_parallel.py +0 -0
  76. {diffct-2.3.0 → diffct-2.4.0}/tests/distributed_projector_check.py +0 -0
  77. {diffct-2.3.0 → diffct-2.4.0}/tests/test_adjoint_inner_product.py +0 -0
  78. {diffct-2.3.0 → diffct-2.4.0}/tests/test_allocator_capacity.py +0 -0
  79. {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_autograd.py +0 -0
  80. {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_distribution.py +0 -0
  81. {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_regressions.py +0 -0
  82. {diffct-2.3.0 → diffct-2.4.0}/tests/test_analytical_weight_distribution.py +0 -0
  83. {diffct-2.3.0 → diffct-2.4.0}/tests/test_angular_low_precision.py +0 -0
  84. {diffct-2.3.0 → diffct-2.4.0}/tests/test_axis_aligned_chords.py +0 -0
  85. {diffct-2.3.0 → diffct-2.4.0}/tests/test_batch_execution_edges.py +0 -0
  86. {diffct-2.3.0 → diffct-2.4.0}/tests/test_batched_surfaces.py +0 -0
  87. {diffct-2.3.0 → diffct-2.4.0}/tests/test_block_cgls.py +0 -0
  88. {diffct-2.3.0 → diffct-2.4.0}/tests/test_block_storage.py +0 -0
  89. {diffct-2.3.0 → diffct-2.4.0}/tests/test_cgls_state_admission.py +0 -0
  90. {diffct-2.3.0 → diffct-2.4.0}/tests/test_completed_transfer_reuse.py +0 -0
  91. {diffct-2.3.0 → diffct-2.4.0}/tests/test_cone_projector_autograd.py +0 -0
  92. {diffct-2.3.0 → diffct-2.4.0}/tests/test_cpu_install.py +0 -0
  93. {diffct-2.3.0 → diffct-2.4.0}/tests/test_cuda_exit_deallocations.py +0 -0
  94. {diffct-2.3.0 → diffct-2.4.0}/tests/test_cuda_sink_lifetimes.py +0 -0
  95. {diffct-2.3.0 → diffct-2.4.0}/tests/test_cuda_smoke.py +0 -0
  96. {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surface_dtypes.py +0 -0
  97. {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surface_mapping.py +0 -0
  98. {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surface_units.py +0 -0
  99. {diffct-2.3.0 → diffct-2.4.0}/tests/test_detector_surfaces.py +0 -0
  100. {diffct-2.3.0 → diffct-2.4.0}/tests/test_example_psnr.py +0 -0
  101. {diffct-2.3.0 → diffct-2.4.0}/tests/test_execution_buffers.py +0 -0
  102. {diffct-2.3.0 → diffct-2.4.0}/tests/test_execution_features.py +0 -0
  103. {diffct-2.3.0 → diffct-2.4.0}/tests/test_execution_schedules.py +0 -0
  104. {diffct-2.3.0 → diffct-2.4.0}/tests/test_fbp_fan_accuracy.py +0 -0
  105. {diffct-2.3.0 → diffct-2.4.0}/tests/test_fbp_fan_offsets.py +0 -0
  106. {diffct-2.3.0 → diffct-2.4.0}/tests/test_fbp_parallel_accuracy.py +0 -0
  107. {diffct-2.3.0 → diffct-2.4.0}/tests/test_fdk_cone_accuracy.py +0 -0
  108. {diffct-2.3.0 → diffct-2.4.0}/tests/test_fdk_cone_offsets.py +0 -0
  109. {diffct-2.3.0 → diffct-2.4.0}/tests/test_geometry_autograd.py +0 -0
  110. {diffct-2.3.0 → diffct-2.4.0}/tests/test_geometry_gradients.py +0 -0
  111. {diffct-2.3.0 → diffct-2.4.0}/tests/test_gradcheck.py +0 -0
  112. {diffct-2.3.0 → diffct-2.4.0}/tests/test_lazy_native_observers.py +0 -0
  113. {diffct-2.3.0 → diffct-2.4.0}/tests/test_multi_gpu_exit.py +0 -0
  114. {diffct-2.3.0 → diffct-2.4.0}/tests/test_native_surface_transforms.py +0 -0
  115. {diffct-2.3.0 → diffct-2.4.0}/tests/test_output_pipeline_overlap.py +0 -0
  116. {diffct-2.3.0 → diffct-2.4.0}/tests/test_partition_lifecycle.py +0 -0
  117. {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_distributed_geometry.py +0 -0
  118. {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_endpoint_precision.py +0 -0
  119. {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_execution_bounds.py +0 -0
  120. {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_numeric_regressions.py +0 -0
  121. {diffct-2.3.0 → diffct-2.4.0}/tests/test_portable_surface_admission.py +0 -0
  122. {diffct-2.3.0 → diffct-2.4.0}/tests/test_projector_api.py +0 -0
  123. {diffct-2.3.0 → diffct-2.4.0}/tests/test_projector_degenerate_geometry.py +0 -0
  124. {diffct-2.3.0 → diffct-2.4.0}/tests/test_projector_geometry_devices.py +0 -0
  125. {diffct-2.3.0 → diffct-2.4.0}/tests/test_ramp_filter_windows.py +0 -0
  126. {diffct-2.3.0 → diffct-2.4.0}/tests/test_reconstruction_api.py +0 -0
  127. {diffct-2.3.0 → diffct-2.4.0}/tests/test_reference_backend.py +0 -0
  128. {diffct-2.3.0 → diffct-2.4.0}/tests/test_release_2_0_1.py +0 -0
  129. {diffct-2.3.0 → diffct-2.4.0}/tests/test_review_230.py +0 -0
  130. {diffct-2.3.0 → diffct-2.4.0}/tests/test_review_regressions.py +0 -0
  131. {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_admission_regressions.py +0 -0
  132. {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_api.py +0 -0
  133. {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_checkpoint_edges.py +0 -0
  134. {diffct-2.3.0 → diffct-2.4.0}/tests/test_solver_collective_failures.py +0 -0
  135. {diffct-2.3.0 → diffct-2.4.0}/tests/test_spatial_partition.py +0 -0
  136. {diffct-2.3.0 → diffct-2.4.0}/tests/test_training_transforms.py +0 -0
  137. {diffct-2.3.0 → diffct-2.4.0}/tests/test_transfer_pipeline.py +0 -0
  138. {diffct-2.3.0 → diffct-2.4.0}/tests/test_transfer_retirement_lifetimes.py +0 -0
  139. {diffct-2.3.0 → diffct-2.4.0}/tests/test_tv_zero_dependency.py +0 -0
  140. {diffct-2.3.0 → diffct-2.4.0}/tests/test_two_far_endpoints.py +0 -0
  141. {diffct-2.3.0 → diffct-2.4.0}/tests/test_weights.py +0 -0
@@ -5,6 +5,32 @@ All notable changes to this project are documented in this file.
5
5
  The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
6
6
  and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
7
7
 
8
+ ## [Unreleased]
9
+
10
+ ## [2.4.0] - 2026-10-20
11
+
12
+ ### Added
13
+
14
+ - Per-call `trajectory` overrides on `Projector` operations, preserving the constructor default and each call's geometry graph without rebuilding the projector (#29).
15
+ - Angle-tensor inputs for circular trajectory helpers and differentiable `transform_trajectory` rigid transforms (#29).
16
+ - `beam="parallel3d"` and `circular_trajectory_3d_parallel`, using infinite parallel rays and the PyTorch backend on CPU and CUDA, with data and first-order geometry gradients (#29).
17
+ - `Projector(native_volume_layout=True)` returns native full-CUDA cone backprojections of unbatched flat-detector inputs as `(D, H, W)` tensors with the kernels' `(W, H, D)`-contiguous strides instead of contiguous copies, and native cone projections accept such tensors without copying. An iterative solver that keeps its volumes in this layout saves two whole-volume copies per normal-operator application on one device. The default is unchanged (#27).
18
+
19
+ ### Changed
20
+
21
+ - The native cone backprojection kernel steps through the z crossings of each xy cell in an inner loop and groups adjacent detector rows in a warp, so the atomic additions of a warp hit neighbouring voxels of one column. Each ray's segments and arithmetic are unchanged; backprojection takes 1.25–1.8x less time on large cone and helical problems, depending on the geometry (#35).
22
+ - Projector execution planning takes no allocator snapshot when a full launch clearly fits: lower bounds of the memory limits come from `mem_get_info` and one statistics query, and the exact planner runs only when they do not suffice. Plans are unchanged. Host time per call for small problems drops by about 1 ms. Distributed operators keep the exact planner (#26).
23
+ - Distributed calls and solvers check rank agreement with one fixed-size collective per request instead of Python object exchanges, skip the callback protocol when no callback is given, and accumulate solver norms on the device with one host read per norm. Results are unchanged (#34).
24
+ - The introduction video explains 2.3.0 tiled execution and view, space and hybrid partitions with a computed walnut example; the README GIF is 12 MB instead of 22 MB (#25).
25
+
26
+ ### Fixed
27
+
28
+ - Detector surfaces expand world points on each compute device after view sharding, preserving native composition precision. Memory planning uses local view counts and composition dtypes. Streamed execution accepts CPU or CUDA offsets from both surface interfaces (#30).
29
+ - Surface snapshots preserve gradient connections on CPU. Streamed composition and geometry backward use bounded fragments. Distributed calls check callback layouts and propagate validation errors before numerical work (#30).
30
+ - Native surface admission sizes validation batches from available memory and staging limits, replacing fixed one-view batches (#32).
31
+ - Rigid trajectory transforms accept valid surface detector frames whose reference center coincides with the source or whose reference plane is edge-on (#29).
32
+ - Call-time trajectory validation checks flat detectors by corners and coordinate bounds, with exact pixel checks for ambiguous views. Frames stage once per validation device. Native surface validation uses bounded batches on each compute device after view sharding (#29).
33
+
8
34
  ## [2.3.0] - 2026-10-14
9
35
 
10
36
  ### Added
@@ -13,8 +13,8 @@ authors:
13
13
  given-names: Yipeng
14
14
  email: yipeng.sun@fau.de
15
15
  alias: sypsyp97
16
- version: 2.3.0
17
- date-released: "2026-10-14"
16
+ version: 2.4.0
17
+ date-released: "2026-10-20"
18
18
  license: Apache-2.0
19
19
  repository-code: "https://github.com/sypsyp97/diffct"
20
20
  url: "https://sypsyp97.github.io/diffct/"
@@ -43,9 +43,9 @@ preferred-citation:
43
43
  - family-names: Maier
44
44
  given-names: Andreas
45
45
  year: 2026
46
- month: 5
46
+ month: 10
47
47
  journal: Preprints
48
48
  publisher:
49
49
  name: Preprints
50
- doi: 10.20944/preprints202605.1446.v1
51
- url: "https://doi.org/10.20944/preprints202605.1446.v1"
50
+ doi: 10.20944/preprints202605.1446.v2
51
+ url: "https://doi.org/10.20944/preprints202605.1446.v2"
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.5
2
2
  Name: diffct
3
- Version: 2.3.0
3
+ Version: 2.4.0
4
4
  Summary: Differentiable CT projection and reconstruction: precision modes, streaming, distributed execution and geometry gradients
5
5
  Project-URL: Homepage, https://github.com/sypsyp97/diffct
6
6
  Project-URL: Documentation, https://sypsyp97.github.io/diffct/
@@ -52,8 +52,8 @@ Description-Content-Type: text/markdown
52
52
 
53
53
  <p align="center">
54
54
  <a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
55
- <a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
56
- <a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
55
+ <a href="https://www.preprints.org/manuscript/202605.1446/v2">Technical report</a> ·
56
+ <a href="https://doi.org/10.20944/preprints202605.1446.v2">DOI</a>
57
57
  </p>
58
58
 
59
59
  <p align="center">
@@ -68,16 +68,18 @@ Description-Content-Type: text/markdown
68
68
  > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
69
69
  > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
70
70
 
71
+ This README describes the source checkout, including [unreleased APIs](CHANGELOG.md#unreleased). Per-call trajectories, angle-tensor circular helpers, `transform_trajectory`, `parallel3d` and `native_volume_layout` require installation from source.
72
+
71
73
  ## Why diffct
72
74
 
73
75
  To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
74
76
 
75
- - **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
77
+ - **Arbitrary trajectories.** Each view has its own ray direction or source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
76
78
  - **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
77
79
  - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
78
80
  - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
79
81
  - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
80
- - **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts and validation limits](docs/EXECUTION.md).
82
+ - **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts](docs/EXECUTION.md).
81
83
  - **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
82
84
  - **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
83
85
 
@@ -99,7 +101,7 @@ From source, with the examples:
99
101
  ```bash
100
102
  git clone https://github.com/sypsyp97/diffct.git
101
103
  cd diffct && pip install -e ".[cu12]"
102
- python examples/quickstart.py # smoke test: prints adjoint mismatch ~1e-8 for each beam
104
+ python examples/quickstart.py # prints adjoint mismatch and gradients for each beam
103
105
  ```
104
106
 
105
107
  <details>
@@ -138,7 +140,11 @@ loss.backward() # x.grad = A^T (A x - y)
138
140
 
139
141
  **Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#custom-or-calibrated-trajectories).
140
142
 
141
- **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
143
+ **Geometry gradients.** Reuse one projector with `A.project(volume, trajectory=current)` or `A.backproject(sinogram, trajectory=current)`. Circular helpers accept learnable angle tensors, and `transform_trajectory` applies differentiable rigid transforms. Omitting the override keeps the constructor trajectory. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
144
+
145
+ **3D parallel beams.** Use `circular_trajectory_3d_parallel(...)` and `beam="parallel3d"`. The PyTorch backend provides matched projection and backprojection on CPU and CUDA.
146
+
147
+ **Native volume layout.** For unbatched, flat-detector cone operations, set `native_volume_layout=True` to avoid volume layout copies. This applies to the native full-CUDA path on one compute device. Shapes stay `(D, H, W)`; the returned strides follow the kernel layout. The default is `False`. See [chunked execution](docs/REFERENCE.md#chunked-execution) for limits.
142
148
 
143
149
  **Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
144
150
 
@@ -171,18 +177,18 @@ volume slabs use `--partition space` under torchrun. Run
171
177
 
172
178
  <p align="center">
173
179
  <picture>
174
- <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
175
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.2.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
180
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.4.0">
181
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
176
182
  </picture>
177
183
  </p>
178
184
 
179
185
  | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
180
186
  |---|---:|---:|---:|
181
- | Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
182
- | Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
183
- | CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
187
+ | Forward projection | 493.40 ms | 131.78 ms (3.7×) | 72.42 ms (6.8×) |
188
+ | Adjoint (backprojection) | 1118.98 ms | 293.01 ms (3.8×) | 232.50 ms (4.8×) |
189
+ | CGLS iteration | 1744.92 ms | 465.62 ms (3.7×) | 337.98 ms (5.2×) |
184
190
 
185
- diffct 2.2.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
191
+ diffct 2.4.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
186
192
 
187
193
  Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
188
194
 
@@ -222,7 +228,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
222
228
 
223
229
  ## Examples
224
230
 
225
- `quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
231
+ `quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `benchmark_trajectory_admission.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
226
232
 
227
233
  ## Documentation
228
234
 
@@ -253,12 +259,12 @@ If you use this library in your research, please cite the software and the techn
253
259
  }
254
260
 
255
261
  @article{202605.1446,
256
- doi = {10.20944/preprints202605.1446.v1},
257
- url = {https://doi.org/10.20944/preprints202605.1446.v1},
262
+ doi = {10.20944/preprints202605.1446.v2},
263
+ url = {https://doi.org/10.20944/preprints202605.1446.v2},
258
264
  year = 2026,
259
- month = {May},
265
+ month = {Oct},
260
266
  publisher = {Preprints},
261
- author = {Yipeng Sun and Linda-Sophie Schneider and Chengze ye and Andreas Maier},
267
+ author = {Yipeng Sun and Linda-Sophie Schneider and Chengze Ye and Andreas Maier},
262
268
  title = {diffct: Differentiable CT Operators from Circular Orbits to Arbitrary Trajectories},
263
269
  journal = {Preprints}
264
270
  }
@@ -21,8 +21,8 @@
21
21
 
22
22
  <p align="center">
23
23
  <a href="https://github.com/sypsyp97/diffct/blob/main/docs/assets/diffct_intro.mp4">Intro video (MP4)</a> ·
24
- <a href="https://www.preprints.org/manuscript/202605.1446/v1">Technical report</a> ·
25
- <a href="https://doi.org/10.20944/preprints202605.1446.v1">DOI</a>
24
+ <a href="https://www.preprints.org/manuscript/202605.1446/v2">Technical report</a> ·
25
+ <a href="https://doi.org/10.20944/preprints202605.1446.v2">DOI</a>
26
26
  </p>
27
27
 
28
28
  <p align="center">
@@ -37,16 +37,18 @@
37
37
  > **Version 2.0** replaces the circular-orbit API of 1.x with per-view trajectories. Code written for 1.x needs changes; see the [migration notes](https://github.com/sypsyp97/diffct/blob/main/docs/MIGRATION.md).
38
38
  > For Apple Silicon, use the [MLX port by Linda-Sophie Schneider](https://github.com/Linda-SophieSchneider/DiffCT-MLX).
39
39
 
40
+ This README describes the source checkout, including [unreleased APIs](CHANGELOG.md#unreleased). Per-call trajectories, angle-tensor circular helpers, `transform_trajectory`, `parallel3d` and `native_volume_layout` require installation from source.
41
+
40
42
  ## Why diffct
41
43
 
42
44
  To our knowledge, diffct is the only open-source GPU CT library that combines arbitrary per-view trajectories, autograd through the volume, first-order gradients through the acquisition geometry, and built-in multi-GPU and multi-node execution. We compared it with LEAP, TIGRE, ASTRA/tomosipo, DiffDRR and others in October 2026.
43
45
 
44
- - **Arbitrary trajectories.** Each view has its own source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
46
+ - **Arbitrary trajectories.** Each view has its own ray direction or source, detector centre and detector axes. Circular, helical, saddle, sinusoidal, random and calibrated scans use the same code. Trajectory tensors with `requires_grad=True` receive geometry gradients for calibration.
45
47
  - **Curved detectors.** `ParameterizedSurface` defines native arc, cylindrical or arbitrary pixels from global coordinates and explicit parameters, sampled in bounded view/pixel batches. Legacy `detector_surface(u, v)` callbacks remain available. Projection and matched backprojection support first-order surface gradients. See [the surface API](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#parameterized-detector-surfaces).
46
48
  - **Matched operators.** `project()` and `backproject()` form an exact adjoint pair for the cell-constant Siddon model. Both support PyTorch autograd, with volume/sinogram gradients and Hessian-vector products.
47
49
  - **Large volumes.** Automatic spatial, view and pixel blocks with bounded transfer buffers control working memory. `project_into`/`backproject_into` accept caller-owned tensor or disk stores; tensor autograd keeps complete data on its input device. See [chunked execution](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#chunked-execution).
48
50
  - **Multi-GPU and multi-node.** Use `devices=[0, 1, 2, 3]` in one process, or torchrun and NCCL across ranks. Views, space and hybrid partitions control replication of volumes and projections. Speedup depends on workload and communication costs.
49
- - **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts and validation limits](docs/EXECUTION.md).
51
+ - **Precision and training.** Explicit native/PyTorch backends, FP32/FP64 modes, per-axis voxel spacing, leading batch/channel dimensions and data `torch.func` transforms. The portable path supports CPU and PyTorch CUDA/ROCm devices; see [execution contracts](docs/EXECUTION.md).
50
52
  - **Solver API.** `cgls`, `sirt` and `tv_reconstruction` report residual history, tolerances and callbacks, with scan/data fingerprints and complete checkpoint state. `cgls_into` keeps state in bounded mapped blocks.
51
53
  - **Differentiable FBP and FDK.** `fbp`, `fdk` and approximate `helical_fdk` compose the public filters, weights and analytical backprojectors. Their backward kernels transpose the actual interpolation stencil; custom filters and weights can be learned. See the [analytical guide](docs/source/analytical_learning.rst).
52
54
 
@@ -68,7 +70,7 @@ From source, with the examples:
68
70
  ```bash
69
71
  git clone https://github.com/sypsyp97/diffct.git
70
72
  cd diffct && pip install -e ".[cu12]"
71
- python examples/quickstart.py # smoke test: prints adjoint mismatch ~1e-8 for each beam
73
+ python examples/quickstart.py # prints adjoint mismatch and gradients for each beam
72
74
  ```
73
75
 
74
76
  <details>
@@ -107,7 +109,11 @@ loss.backward() # x.grad = A^T (A x - y)
107
109
 
108
110
  **Custom or calibrated trajectories.** Pass calibrated tensors directly instead of a generator; no circular fit is needed. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#custom-or-calibrated-trajectories).
109
111
 
110
- **Geometry gradients.** Set `requires_grad=True` on trajectory tensors before you build `Projector`. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
112
+ **Geometry gradients.** Reuse one projector with `A.project(volume, trajectory=current)` or `A.backproject(sinogram, trajectory=current)`. Circular helpers accept learnable angle tensors, and `transform_trajectory` applies differentiable rigid transforms. Omitting the override keeps the constructor trajectory. See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#geometry-gradients).
113
+
114
+ **3D parallel beams.** Use `circular_trajectory_3d_parallel(...)` and `beam="parallel3d"`. The PyTorch backend provides matched projection and backprojection on CPU and CUDA.
115
+
116
+ **Native volume layout.** For unbatched, flat-detector cone operations, set `native_volume_layout=True` to avoid volume layout copies. This applies to the native full-CUDA path on one compute device. Shapes stay `(D, H, W)`; the returned strides follow the kernel layout. The default is `False`. See [chunked execution](docs/REFERENCE.md#chunked-execution) for limits.
111
117
 
112
118
  **Several GPUs and nodes.** See [REFERENCE](https://github.com/sypsyp97/diffct/blob/main/docs/REFERENCE.md#several-gpus-and-nodes) and [DISTRIBUTED.md](https://github.com/sypsyp97/diffct/blob/main/docs/DISTRIBUTED.md).
113
119
 
@@ -140,18 +146,18 @@ volume slabs use `--partition space` under torchrun. Run
140
146
 
141
147
  <p align="center">
142
148
  <picture>
143
- <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.2.0">
144
- <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.2.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
149
+ <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_dark.png?v=2.4.0">
150
+ <img src="https://raw.githubusercontent.com/sypsyp97/diffct/main/docs/assets/scaling_light.png?v=2.4.0" width="100%" alt="Time per forward, adjoint and CGLS iteration on 1, 4 and 8 GPUs, 512^3 volume">
145
151
  </picture>
146
152
  </p>
147
153
 
148
154
  | 512³ volume | 1 GPU | 4 GPUs, 1 node | 8 GPUs, 2 nodes |
149
155
  |---|---:|---:|---:|
150
- | Forward projection | 493.05 ms | 132.22 ms (3.7×) | 73.18 ms (6.7×) |
151
- | Adjoint (backprojection) | 1400.80 ms | 369.29 ms (3.8×) | 265.87 ms (5.3×) |
152
- | CGLS iteration | 2054.36 ms | 547.06 ms (3.8×) | 373.38 ms (5.5×) |
156
+ | Forward projection | 493.40 ms | 131.78 ms (3.7×) | 72.42 ms (6.8×) |
157
+ | Adjoint (backprojection) | 1118.98 ms | 293.01 ms (3.8×) | 232.50 ms (4.8×) |
158
+ | CGLS iteration | 1744.92 ms | 465.62 ms (3.7×) | 337.98 ms (5.2×) |
153
159
 
154
- diffct 2.2.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
160
+ diffct 2.4.0, FP32, A100 64 GB GPUs; circular cone beam, 360 views, 1024×768 detector, pitch 1.25. Default memory settings, without an imposed allocator cap.
155
161
 
156
162
  Projection times are medians of seven samples after three warmups. CGLS uses three ten-iteration solves after a two-iteration warmup; the median includes initialization and is divided by ten. Timings use synchronized wall time and the slowest rank. Raw samples and configuration: [scaling.json](https://github.com/sypsyp97/diffct/blob/main/docs/assets/scaling.json).
157
163
 
@@ -191,7 +197,7 @@ Walnut data: Meaney 2022, Zenodo 6986012, CC BY 4.0; see [examples/data/NOTICE](
191
197
 
192
198
  ## Examples
193
199
 
194
- `quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
200
+ `quickstart.py`, `curved_detector.py`, `chunked_reconstruction.py`, `disk_reconstruction.py`, `benchmark_execution.py`, `benchmark_trajectory_admission.py`, `analytical_reconstruction.py`, `learnable_reconstruction.py`, `helical_fdk.py`, `iterative_reconstruction.py`, `walnut_reconstruction.py`, `geometry_calibration.py`, `benchmark_projector.py`, `benchmark_scaling.py`, `plot_trajectory.py`, and the Slurm template `slurm/multi_node.sbatch`. Launch modes and distributed-loss rules are in [examples/README.md](https://github.com/sypsyp97/diffct/blob/main/examples/README.md).
195
201
 
196
202
  ## Documentation
197
203
 
@@ -222,12 +228,12 @@ If you use this library in your research, please cite the software and the techn
222
228
  }
223
229
 
224
230
  @article{202605.1446,
225
- doi = {10.20944/preprints202605.1446.v1},
226
- url = {https://doi.org/10.20944/preprints202605.1446.v1},
231
+ doi = {10.20944/preprints202605.1446.v2},
232
+ url = {https://doi.org/10.20944/preprints202605.1446.v2},
227
233
  year = 2026,
228
- month = {May},
234
+ month = {Oct},
229
235
  publisher = {Preprints},
230
- author = {Yipeng Sun and Linda-Sophie Schneider and Chengze ye and Andreas Maier},
236
+ author = {Yipeng Sun and Linda-Sophie Schneider and Chengze Ye and Andreas Maier},
231
237
  title = {diffct: Differentiable CT Operators from Circular Orbits to Arbitrary Trajectories},
232
238
  journal = {Preprints}
233
239
  }
@@ -30,6 +30,8 @@ from .reconstruction import (
30
30
  from .solvers import SolverResult, cgls, sirt, tv_reconstruction, cgls_into, total_variation
31
31
 
32
32
  from .geometry import (
33
+ transform_trajectory,
34
+ circular_trajectory_3d_parallel,
33
35
  circular_trajectory_3d,
34
36
  random_trajectory_3d,
35
37
  spiral_trajectory_3d,
@@ -62,7 +64,7 @@ try:
62
64
  except ImportError:
63
65
  pass
64
66
 
65
- __version__ = '2.3.0'
67
+ __version__ = '2.4.0'
66
68
 
67
69
  __all__ = [
68
70
  'Projector',
@@ -88,6 +90,8 @@ __all__ = [
88
90
  'FanBackprojectorFunction',
89
91
  'ConeProjectorFunction',
90
92
  'ConeBackprojectorFunction',
93
+ 'transform_trajectory',
94
+ 'circular_trajectory_3d_parallel',
91
95
  'circular_trajectory_3d',
92
96
  'random_trajectory_3d',
93
97
  'spiral_trajectory_3d',
@@ -11,6 +11,7 @@ import torch.distributed as dist
11
11
 
12
12
  from .pipeline import _fragments
13
13
  from .storage import NpyStore, TensorStore
14
+ from .distribution import control_status, preflight
14
15
 
15
16
 
16
17
  def count(value, name, positive=False):
@@ -25,47 +26,56 @@ def finite(value, name, positive=False):
25
26
  return float(value)
26
27
 
27
28
 
28
- def agreement(operator, error=None, token=None):
29
+ def agreement(operator, error=None, token=None, *, cache_key=None):
29
30
  """Every rank rejects a bad request before numerical execution starts."""
30
31
  world = getattr(operator, "world_size", 1)
31
32
  if world == 1:
32
33
  if error is not None:
33
34
  raise error
34
35
  return
35
- records = [None] * world
36
- dist.all_gather_object(records, (None if error is None else str(error), token),
37
- group=operator.process_group)
38
- failures = [record[0] for record in records if record[0] is not None]
39
- if failures or any(record[1] != records[0][1] for record in records):
40
- raise ValueError("distributed solver agreement failed: " + (failures[0] if failures else "inconsistent checkpoint iteration or controls"))
36
+ error = None if error is None else ValueError("distributed solver agreement failed: " + str(error))
37
+ mismatch = "distributed solver agreement failed: inconsistent checkpoint iteration or controls"
38
+ partition = getattr(operator, "_partition", None)
39
+ if partition is not None:
40
+ partition.preflight(error, token, cache_key=cache_key, mismatch=mismatch)
41
+ else:
42
+ preflight(error, token, process_group=operator.process_group, world_size=world,
43
+ mismatch=mismatch)
41
44
 
42
45
 
43
- def scalar_sum(operator, value, kind):
44
- """Count physical rays/voxels once, including view/space replicas."""
45
- world = getattr(operator, "world_size", 1)
46
+ def owns_values(operator, kind):
47
+ """Select one copy of each physical ray or voxel for scalar reductions."""
46
48
  partition = getattr(operator, "partition", "views")
47
49
  rank = getattr(operator, "rank", 0)
48
- owns = (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
50
+ return (getattr(operator, "view_rank", 0) == 0 if kind == "volume"
49
51
  else getattr(operator, "space_rank", 0) == 0) if partition == "hybrid" else (
50
52
  rank == 0 if (kind == "volume") == (partition == "views") else True
51
53
  )
52
- value = float(value) if owns else 0.
54
+
55
+
56
+ def scalar_sum(operator, value, kind):
57
+ """Count physical rays/voxels once, including view/space replicas."""
58
+ world = getattr(operator, "world_size", 1)
59
+ value = value if owns_values(operator, kind) else 0.
53
60
  if world > 1:
54
61
  group = operator.process_group
55
62
  device = torch.device("cuda", torch.cuda.current_device()) if dist.get_backend(group) == "nccl" else torch.device("cpu")
56
- accumulator = torch.tensor(value, dtype=torch.float64, device=device)
63
+ accumulator = torch.as_tensor(value, dtype=torch.float64, device=device).clone()
57
64
  dist.all_reduce(accumulator, group=group)
58
- value = accumulator.item()
59
- return value
65
+ return accumulator.item()
66
+ return float(value)
60
67
 
61
68
 
62
69
  def squared_norm(operator, value, kind, elements=1048576):
63
70
  store = TensorStore(value) if isinstance(value, torch.Tensor) else value
64
71
  total = 0.
65
- for selected in _fragments(tuple(store.shape), elements):
66
- # Clone even an FP64 view: norms must never alter caller measurements.
67
- block = store.read(selected).to(dtype=torch.float64).clone()
68
- total += block.square_().sum().item()
72
+ if owns_values(operator, kind):
73
+ total = torch.zeros((), dtype=torch.float64, device=getattr(store, "device", "cpu"))
74
+ for selected in _fragments(tuple(store.shape), elements):
75
+ # Clone even an FP64 view: norms must never alter caller measurements.
76
+ block = store.read(selected).to(dtype=torch.float64).clone()
77
+ # Preserve the block reductions and their sequential FP64 order.
78
+ total.add_(block.square_().sum().to(device=total.device))
69
79
  return scalar_sum(operator, total, kind)
70
80
 
71
81
 
@@ -270,15 +280,22 @@ def save_checkpoint(operator, directory, states, metadata, iteration, history, m
270
280
  agreement(operator, error, iteration)
271
281
 
272
282
 
283
+ def callbacks_enabled(operator, callback):
284
+ """Check callback presence once, including callbacks on only some ranks."""
285
+ return bool(control_status((int(callback is not None),),
286
+ process_group=getattr(operator, "process_group", None),
287
+ world_size=getattr(operator, "world_size", 1))[0])
288
+
289
+
273
290
  def callback_stop(operator, callback, iteration, residual, volume):
274
291
  error, stop = None, False
275
292
  try:
276
293
  stop = callback is not None and callback(iteration, residual, volume) is False
277
294
  except Exception as caught:
278
295
  error = caught
279
- agreement(operator, error)
280
- if getattr(operator, "world_size", 1) > 1:
281
- records = [None] * operator.world_size
282
- dist.all_gather_object(records, stop, group=operator.process_group)
283
- stop = any(records)
284
- return stop
296
+ failed, stop = control_status((int(error is not None), int(stop)),
297
+ process_group=getattr(operator, "process_group", None),
298
+ world_size=getattr(operator, "world_size", 1))
299
+ if failed:
300
+ agreement(operator, error)
301
+ return bool(stop)