[BUG] `parallel=True` in `TriangleScene.compute_paths` not working with `jax>=0.6`
bughelp wanted
### Terms
- [x] Checked the [existing issues](https://github.com/jeertmans/differt/issues?q=is%3Aissue+label%3Abug+) and [discussions](https://github.com/jeertmans/differt/discussions) to see if my issue had not already been reported;
- [x] Checked the [documentation](https://differt.rtfd.io/);
- [x] Read the [installation instructions](https://differt.rtfd.io/latest/installation.html);
- [x] Created a virtual environment in which I can reproduce my bug;
### Describe the issue
As observed in CI when JAX v0.6 was just released, tests using manual shardings (for parallel computation) were failing. Apparently, this bug can be fixed using `jax.lax.pvary` (see https://docs.jax.dev/en/latest/notebooks/shard_map.html#scan-vma) at places where it is recommended by the errors.
While this works, I am not 100% happy with this solution, for two reasons:
1. `jax.lax.pvary` is currently not documented (or I couldn't find it in the docs);
2. using `jax.lax.pvary` looks very constraining to me, as we need to know the name of each sharding axis (e.g., `'i'` or `'j'`), which eventually depends on the upstream caller. Hard-coding those names will probably limit the use of other axis names or shardings in general, which I am not really a fan of.
I am planning on disabling `parallel` computation on `jax>=0.6` until we find a better fix. If you would like to contribute to this issue, please feel free to reach out!
### Command
```shell
uv run pytest -k parallel -n0 -x
```
### Issue Type
Error or warning when executing the code
### Python version
Python 3.11.10
### Python environment
```shell
absl-py==2.3.0
accessible-pygments==0.0.5
alabaster==1.0.0
annotated-types==0.7.0
anyio==4.9.0
argon2-cffi==25.1.0
argon2-cffi-bindings==21.2.0
arrow==1.3.0
asttokens==3.0.0
async-lru==2.0.5
attrs==25.3.0
babel==2.17.0
beartype==0.21.0
beautifulsoup4==4.13.4
bleach==6.2.0
bracex==2.6
bump-my-version==1.2.0
certifi==2025.6.15
cffi==1.17.1
cfgv==3.4.0
charset-normalizer==3.4.2
chex==0.1.90
click==8.2.1
comm==0.2.2
contourpy==1.3.2
coverage==7.9.1
cycler==0.12.1
debugpy==1.8.14
decorator==5.2.1
defusedxml==0.7.1
-e file:///export/home/eertmans/repositories/DiffeRT/DiffeRT/differt
-e file:///export/home/eertmans/repositories/DiffeRT/DiffeRT/differt-core
distlib==0.3.9
docutils==0.21.2
drjit==1.0.3
e3x==1.0.2
equinox==0.12.2
etils==1.12.2
execnet==2.1.1
executing==2.2.0
fastjsonschema==2.21.1
filelock==3.18.0
flax==0.10.6
fonttools==4.58.4
fqdn==1.5.1
freetype-py==2.5.1
fsspec==2025.5.1
greenlet==3.2.3
h11==0.16.0
hsluv==5.0.4
httpcore==1.0.9
httpx==0.28.1
humanize==4.12.3
identify==2.6.12
idna==3.10
imagesize==1.4.1
importlib-metadata==8.7.0
importlib-resources==6.5.2
iniconfig==2.1.0
ipydatawidgets==4.3.5
ipykernel==6.29.5
ipympl==0.9.7
ipython==9.3.0
ipython-pygments-lexers==1.1.1
ipywidgets==8.1.7
isoduration==20.11.0
jax==0.6.2
jax-cuda12-pjrt==0.6.2
jax-cuda12-plugin==0.6.2
jaxlib==0.6.2
jaxtyping==0.3.2
jedi==0.19.2
jinja2==3.1.6
json5==0.12.0
jsonpointer==3.0.0
jsonschema==4.24.0
jsonschema-specifications==2025.4.1
jupyter-cache==1.0.1
jupyter-client==8.6.3
jupyter-core==5.8.1
jupyter-events==0.12.0
jupyter-lsp==2.2.5
jupyter-rfb==0.5.3
jupyter-server==2.16.0
jupyter-server-terminals==0.5.3
jupyterlab==4.4.3
jupyterlab-pygments==0.3.0
jupyterlab-server==2.27.3
jupyterlab-widgets==3.0.15
kiwisolver==1.4.8
latexcodec==3.0.1
markdown-it-py==3.0.0
markupsafe==3.0.2
matplotlib==3.10.3
matplotlib-inline==0.1.7
maturin-import-hook==0.3.0
mdit-py-plugins==0.4.2
mdurl==0.1.2
mistune==3.1.3
mitsuba==3.6.2
ml-dtypes==0.5.1
more-itertools==10.7.0
mpmath==1.3.0
msgpack==1.1.1
myst-nb==1.2.0
myst-parser==4.0.1
nbclient==0.10.2
nbconvert==7.16.6
nbformat==5.10.4
nest-asyncio==1.6.0
nodeenv==1.9.1
notebook==7.4.3
notebook-shim==0.2.4
numpy==2.3.1
nvidia-cublas-cu12==12.9.1.4
nvidia-cuda-cupti-cu12==12.9.79
nvidia-cuda-nvcc-cu12==12.9.86
nvidia-cuda-nvrtc-cu12==12.9.86
nvidia-cuda-runtime-cu12==12.9.79
nvidia-cudnn-cu12==9.10.2.21
nvidia-cufft-cu12==11.4.1.4
nvidia-cusolver-cu12==11.7.5.82
nvidia-cusparse-cu12==12.5.10.65
nvidia-nccl-cu12==2.27.3
nvidia-nvjitlink-cu12==12.9.86
nvidia-nvshmem-cu12==3.2.5
opt-einsum==3.4.0
optax==0.2.5
orbax-checkpoint==0.11.16
overrides==7.7.0
packaging==25.0
pandas==2.3.0
pandocfilters==1.5.1
parso==0.8.4
pbr==6.1.1
pexpect==4.9.0
pillow==11.2.1
platformdirs==4.3.8
plotly==5.24.1
pluggy==1.6.0
pre-commit==4.2.0
prometheus-client==0.22.1
prompt-toolkit==3.0.51
protobuf==6.31.1
psutil==7.0.0
ptyprocess==0.7.0
pure-eval==0.2.3
py-cpuinfo==9.0.0
pybtex==0.25.0
pybtex-docutils==1.0.3
pycparser==2.22
pydantic==2.11.7
pydantic-core==2.33.2
pydantic-settings==2.10.1
pydata-sphinx-theme==0.15.4
pygments==2.19.2
pyparsing==3.2.3
pyqt6==6.9.1
pyqt6-qt6==6.9.1
pyqt6-sip==13.10.2
pytest==8.3.4
pytest-benchmark==5.1.0
pytest-codspeed==3.2.0
pytest-cov==6.2.1
pytest-env==1.1.5
pytest-missing-modules==0.2.1
pytest-subtests==0.14.2
pytest-xdist==3.7.0
python-dateutil==2.9.0.post0
python-dotenv==1.1.1
python-json-logger==3.3.0
pythreejs==2.4.2
pytz==2025.2
pyyaml==6.0.2
pyzmq==27.0.0
questionary==2.1.0
referencing==0.36.2
requests==2.32.4
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rich==14.0.0
rich-click==1.8.9
roman-numerals-py==3.1.0
rpds-py==0.25.1
scipy==1.16.0
send2trash==1.8.3
setuptools==80.9.0
simplejson==3.20.1
sionna-rt==1.1.0
six==1.17.0
sniffio==1.3.1
snowballstemmer==3.0.1
soupsieve==2.7
sphinx==8.2.3
sphinx-autodoc-typehints==3.2.0
sphinx-book-theme==1.1.4
sphinx-copybutton==0.5.2
sphinx-design==0.6.1
sphinx-plotly-directive==0.1.3
sphinx-remove-toctrees==1.0.0.post1
sphinxcontrib-apidoc==0.6.0
sphinxcontrib-applehelp==2.0.0
sphinxcontrib-bibtex==2.6.4
sphinxcontrib-devhelp==2.0.0
sphinxcontrib-htmlhelp==2.1.0
sphinxcontrib-jsmath==1.0.1
sphinxcontrib-qthelp==2.0.0
sphinxcontrib-serializinghtml==2.0.0
sphinxext-opengraph==0.10.0
sqlalchemy==2.0.41
stack-data==0.6.3
sympy==1.14.0
tabulate==0.9.0
tenacity==9.1.2
tensorstore==0.1.75
terminado==0.18.1
tinycss2==1.4.0
tomlkit==0.13.3
toolz==1.0.0
tornado==6.5.1
tqdm==4.67.1
traitlets==5.14.3
traittypes==0.2.1
treescope==0.1.9
types-python-dateutil==2.9.0.20250516
typing-extensions==4.14.0
typing-inspection==0.4.1
tzdata==2025.2
uri-template==1.3.0
urllib3==2.5.0
virtualenv==20.31.2
vispy==0.15.2
wadler-lindig==0.1.7
wcmatch==10.1
wcwidth==0.2.13
webcolors==24.11.1
webencodings==0.5.1
websocket-client==1.8.0
widgetsnbextension==4.0.14
zipp==3.23.0
```
### What is your platform?
Linux
### Other platform
_No response_
### Python code
```python
```
### Relevant terminal output
```shell
================================================================= FAILURES =================================================================
______________________________ TestTriangleScene.test_compute_paths_parallel[exhaustive-8-8-1-1-expectation0] ______________________________
self = <tests.scene.test_triangle_scene.TestTriangleScene object at 0x7ff29033ad10>, m_tx = 8, n_tx = 8, m_rx = 1, n_rx = 1
method = 'exhaustive', expectation = <contextlib.nullcontext object at 0x7ff2d81c77d0>
advanced_path_tracing_example_scene = TriangleScene(
transmitters=f32[3],
receivers=f32[3],
mesh=TriangleMesh(vertices=f32[56,3], triangles=i32[24,3], material_names=())
)
@skip_if_not_8_devices
@pytest.mark.parametrize(
("m_tx", "n_tx", "m_rx", "n_rx", "expectation"),
[
(8, 8, 1, 1, does_not_raise()),
(1, 1, 8, 8, does_not_raise()),
(4, 2, 1, 1, does_not_raise()),
(1, 1, 2, 4, does_not_raise()),
(1, 4, 2, 1, does_not_raise()),
(1, 2, 4, 1, does_not_raise()),
(
7,
1,
1,
1,
pytest.raises(ValueError, match="Found 8 devices available"),
),
(
1,
2,
3,
1,
pytest.raises(ValueError, match="Found 8 devices available"),
),
],
)
@pytest.mark.parametrize(
"method",
[
"exhaustive",
"sbr",
pytest.param(
"hybrid", marks=pytest.mark.xfail(reason="Not yet implemented.")
),
],
)
def test_compute_paths_parallel(
self,
m_tx: int,
n_tx: int,
m_rx: int,
n_rx: int,
method: Literal["exhaustive", "sbr", "hybrid"],
expectation: AbstractContextManager[Exception],
advanced_path_tracing_example_scene: TriangleScene,
) -> None:
scene = advanced_path_tracing_example_scene
scene = scene.with_transmitters_grid(m_tx, n_tx)
scene = scene.with_receivers_grid(m_rx, n_rx)
num_rays = m_rx * n_rx
with expectation:
> paths = scene.compute_paths(
order=1, method=method, num_rays=num_rays, parallel=True
)
differt/tests/scene/test_triangle_scene.py:498:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
differt/src/differt/scene/_triangle_scene.py:1127: in compute_paths
return _compute_paths(
differt/src/differt/scene/_triangle_scene.py:293: in _compute_paths
vertices, mask = fun(tx_vertices, rx_vertices)
differt/src/differt/scene/_triangle_scene.py:129: in fun
paths = image_method(
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
from_vertices = Traced<float32[8,1,24,3]{V:i}>with<DynamicJaxprTrace>, to_vertices = Traced<float32[8,1,24,3]{V:j}>with<DynamicJaxprTrace>
mirror_vertices = Traced<float32[1,24,3]>with<DynamicJaxprTrace>, mirror_normals = Traced<float32[1,24,3]>with<DynamicJaxprTrace>
@jax.jit
def image_method(
from_vertices: Float[ArrayLike, "*#batch 3"],
to_vertices: Float[ArrayLike, "*#batch 3"],
mirror_vertices: Float[ArrayLike, "*#batch num_mirrors 3"],
mirror_normals: Float[ArrayLike, "*#batch num_mirrors 3"],
) -> Float[Array, "*batch num_mirrors 3"]:
"""
Return the ray paths between pairs of vertices, that reflect on a given list of mirrors in between.
The Image Method is a very simple but effective path tracing technique
that can rapidly compute a ray path undergoing a series
of specular reflections on a pre-defined list of mirrors.
The method assumes infinitely long mirrors, and will return invalid
paths in some degenerated cases such as consecutive colinear mirrors,
or impossible configurations. It is the user's responsibility to make
sure that the returned path is correct.
Otherwise, the returned path will, for each reflection, have equal angles
of incidence and of reflection.
Warning:
NaNs and infinity values should be treated as invalid paths, and will naturally
occur when image paths are impossible to trace, e.g., when a mirror
is parallel to a ray segment that it is supposed to reflect.
Args:
from_vertices: An array of ``from`` vertices, i.e., vertices from which the
ray paths start. In a radio communications context, this is usually
an array of transmitters.
to_vertices: An array of ``to`` vertices, i.e., vertices to which the
ray paths end. In a radio communications context, this is usually
an array of receivers.
mirror_vertices: An array of mirror vertices. For each mirror, any
vertex on the infinite plane that describes the mirror is considered
to be a valid vertex.
mirror_normals: An array of mirror normals, where each normal has a unit
length and if perpendicular to the corresponding mirror.
Returns:
An array of ray paths obtained with the image method.
.. note::
The paths do not contain the starting and ending vertices.
You can easily create the complete ray paths using
:func:`assemble_paths<differt.geometry.assemble_paths>`:
.. code-block:: python
paths = image_method(
from_vertices,
to_vertices,
mirror_vertices,
mirror_normals,
)
full_paths = assemble_paths(
from_vertices[..., None, :],
paths,
to_vertices[..., None, :],
)
Examples:
The following image shows how the Image Method (IM) can be applied
to find a path between two nodes (i.e., BS and UE).
.. figure:: ../../_static/image-method.svg
:width: 70%
:align: center
:alt: Image Method example.
Example application of IM in RT. The method determines the only
valid path that can be taken to join BS and UE with, in between, reflection with
two mirrors (the interaction order is important). First, the consecutive images
of the BS are determined through each mirror, using line symmetry. Second,
intersections with mirrors are computed backward, `i.e.`, from last mirror to
first, by joining the UE, then the intersections points, with the images of the
BS. Finally, the valid path can be obtained by joining BS, the intermediary
intersection points, and the UE :cite:`mpt-eucap2023{fig. 5, p. 3}`.
Next, we show how to reproduce the above results using :func:`image_method`.
.. plotly::
>>> from differt.geometry import TriangleMesh, normalize, assemble_paths
>>> from differt.plotting import draw_markers, draw_paths, reuse
>>> from differt.rt import image_method
>>>
>>> from_vertex = jnp.array([+2.0, -1.0, +0.0])
>>> to_vertex = jnp.array([+2.0, +4.0, +0.0])
>>> mirror_vertices = jnp.array([
... [3.0, 3.0, 0.0],
... [4.0, 3.4, 0.0],
... ])
>>> mirror_normals = jnp.array([
... [+1.0, -1.0, +0.0],
... [-1.0, +0.0, +0.0],
... ])
>>> mirror_normals, _ = normalize(mirror_normals)
>>> path = image_method(
... from_vertex,
... to_vertex,
... mirror_vertices,
... mirror_normals,
... )
>>> with reuse(backend="plotly") as fig: # doctest: +SKIP
... TriangleMesh.plane(
... mirror_vertices[0], normal=mirror_normals[0], rotate=-0.954
... ).plot(color="red")
... TriangleMesh.plane(
... mirror_vertices[1], normal=mirror_normals[1]
... ).plot(color="red")
...
... full_path = assemble_paths(
... from_vertex[None, :],
... path,
... to_vertex[None, :],
... )
... draw_paths(
... full_path,
... mode="lines+markers",
... marker={"color": "green"},
... name="Final path",
... )
... markers = jnp.vstack((from_vertex, to_vertex))
... draw_markers(
... markers,
... labels=["BS", "UE"],
... marker={"color": "black"},
... name="BS/UE",
... )
... fig.update_layout(scene_aspectmode="data")
>>> fig # doctest: +SKIP
"""
from_vertices = jnp.asarray(from_vertices)
to_vertices = jnp.asarray(to_vertices)
mirror_vertices = jnp.asarray(mirror_vertices)
mirror_normals = jnp.asarray(mirror_normals)
# Put 'num_mirrors' axis as leading axis
mirror_vertices = jnp.moveaxis(mirror_vertices, -2, 0)
mirror_normals = jnp.moveaxis(mirror_normals, -2, 0)
# Broadcast scan carries, because shapes cannot differ between input and output
batch_and_3 = jnp.broadcast_shapes(
from_vertices.shape,
to_vertices.shape,
mirror_vertices.shape[1:],
mirror_normals.shape[1:],
)
from_vertices = jnp.broadcast_to(from_vertices, batch_and_3)
to_vertices = jnp.broadcast_to(to_vertices, batch_and_3)
def forward(
previous_images: Float[Array, "*batch 3"],
mirror_vertices_and_normals: tuple[
Float[Array, "*#batch 3"], Float[Array, "*#batch 3"]
],
) -> tuple[Float[Array, "*batch 3"], Float[Array, "*batch 3"]]:
# ruff: noqa: DOC201
"""Perform forward pass on vertices by computing consecutive images."""
mirror_vertices, mirror_normals = mirror_vertices_and_normals
images = image_of_vertices_with_respect_to_mirrors(
previous_images,
mirror_vertices,
mirror_normals,
)
return images, images
def backward(
previous_intersections: Float[Array, "*batch 3"],
mirror_vertices_normals_and_images: tuple[
Float[Array, "*#batch 3"],
Float[Array, "*#batch 3"],
Float[Array, "*#batch 3"],
],
) -> tuple[Float[Array, "*batch 3"], Float[Array, "*batch 3"]]:
# ruff: noqa: DOC201
"""Perform backward pass on images by computing the intersection with mirrors."""
mirror_vertices, mirror_normals, images = mirror_vertices_normals_and_images
# We avoid NaNs (caused by subtraction of two infinities) by replacing
# previous_intersections with zeros when they are infinite.
no_previous_intersections = jnp.isinf(previous_intersections)
previous_intersections = jnp.where(
no_previous_intersections,
jnp.zeros_like(previous_intersections),
previous_intersections,
)
intersections = intersection_of_rays_with_planes(
previous_intersections,
images - previous_intersections,
mirror_vertices,
mirror_normals,
)
intersections: Array = jnp.where(
no_previous_intersections,
jnp.full_like(intersections, jnp.inf),
intersections,
)
return intersections, intersections
_, images = jax.lax.scan(
forward,
init=from_vertices,
xs=(mirror_vertices, mirror_normals),
unroll=True,
)
> _, paths = jax.lax.scan(
backward,
init=to_vertices,
xs=(mirror_vertices, mirror_normals, images),
reverse=True,
unroll=True,
)
E TypeError: scan body function carry input and carry output must have equal types, but they differ:
E
E The input carry previous_intersections has type float32[8,1,24,3]{V:j} but the corresponding output carry component has type float32[8,1,24,3]{V:(i,j)}, so the varying manual axes do not match.
E
E This might be fixed by applying `jax.lax.pvary(..., ('i',))` to the initial carry value corresponding to the input carry previous_intersections.
E See https://docs.jax.dev/en/latest/notebooks/shard_map.html#scan-vma for more information.
E
E Revise the function so that all output types match the corresponding input types.
E --------------------
E For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.
differt/src/differt/rt/_image_method.py:351: TypeError
```
### Screenshots
_No response_
### Additional information
Relevant documentation pages:
- https://docs.jax.dev/en/latest/notebooks/shard_map.html
- https://docs.jax.dev/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html
### Recommended fix or suggestions
_No response_
关闭于 2025-07-31 0 条评论