Improve linop.xray.XRayTransform3D

[bwohlberg]

Improvements to linop.xray.XRayTransform3D, in particular:

• Forward and adjoint operations could previously not be jitted.
• Now ≈5x faster than astra rather than ≈10x slower.

Notes:

• Improvements implemented with the assistance of Gemini
• This branch incorporates the fixes from Initialization of scico.linop.xray.XRayTransform3D objects fails #662.
• The approach in this branch includes removal of the manual batching (for memory use reduction) approach in the original implementation, relying instead on jax tools for this purpose, including using batch_size in jax.lax.map, and jax.remat.
• The original manual batching approach is retained in branch brendt/xray-loops. While the implementation in this branch is cleaner, that branch does seem to be ≈10% faster. Further tests (including quantification of both relative computation times and memory usage) and discussion would be appropriate before deciding which solution to adopt.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 93.61%. Comparing base (6e1892e) to head (e4c78a0).

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #665      +/-   ##
==========================================
- Coverage   93.61%   93.61%   -0.00%     
==========================================
  Files          93       93              
  Lines        6431     6429       -2     
==========================================
- Hits         6020     6018       -2     
  Misses        411      411              

Flag	Coverage Δ
unittests	93.61% <100.00%> (-<0.01%)	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.