No error message but failing zero-copy?

[Cattaneo123]

Description

Greetings.
I’m (still) trying to get JAX and PETSc to work together in a no-copy system using the DLPack tools in both. Unfortunately I can’t seem to get it to work right. Ideally, I’d like to create a PETSc vec object using petsc4py, pass it to to a JAX object without copying, make a change to it in a JAX jitted function and have that change reflected in the PETSc object. All of this without copying.

Of note: When I try to do this I get an error that the alignment is wrong and a copy must be made when I call the from-dlpack function, but changing the alignment in the PETSc ./config stage to 32 causes the error message to disappear, even so it still doesn’t function correctly. I’ve tried looking through the documentation, but I’m getting a little turned around.

I’ve included a code snippet below:

from petsc4py import PETSc as PETSc

import jax

from functools import partial

import jax.numpy as jnp

jax.config.update("jax_enable_x64", True) 

@partial(jax.jit, donate_argnums=(0,))

def set_in_place(x):

    return [x.at](http://x.at/)[:].set(3.0)

 

print('\nTesting jax from_dlpack given a PETSc vector that was allocated by PETSc')

x = jnp.ones((1000,1))

y_petsc = PETSc.Vec().createSeq(x.shape[0])

y_petsc.set(0.0)

print(hex(y_petsc.handle))

y2_petsc = PETSc.Vec().createWithDLPack(y_petsc.toDLPack('rw'))

y2_petsc.set(-1.0)

assert y_petsc.getValue(0) == y2_petsc.getValue(0)

print('After creating a second PETSc vector via a DLPack of the first, modifying the memory of one affects the other.')

#y = jnp.from_dlpack(y_petsc.toDLPack('rw'), copy=False)

y = jnp.from_dlpack(y_petsc, copy=False)

orig_ptr = y.unsafe_buffer_pointer()

print(f'before: ptr at {hex(orig_ptr)}')

y = set_in_place(y)

print(f'after:  ptr at {hex(y.unsafe_buffer_pointer())}')

assert orig_ptr == y.unsafe_buffer_pointer()

#assert y_petsc.getValue(0) == y[0], f'The PETSc value {y_petsc.getValue(0)} did not match the JAX value {y[0]}, so modifying the JAX memory did not affect the PETSc memory.'

I’d like the bottom two asserts to pass, but I can only get one of them. I asked about it in discussion a few weeks ago but heard nothing so I thought it may ask through a different avenue, I hope that's alright. I looked into the documentation I couldn't see anything that mentioned why this might be. Furthermore no error is being thrown up either by PETSc or JAX so I'm unsure where the problem may lie.

Alberto

System info (python version, jaxlib version, accelerator, etc.)

jax: 0.6.0
jaxlib: 0.6.0
numpy: 1.26.2
python: 3.11.4 (tags/v3.11.4:d2340ef, Jun 7 2023, 05:45:37) [MSC v.1934 64 bit (AMD64)]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Windows', node='Alberto-Grad', release='10', version='10.0.26100', machine='AMD64')

$ nvidia-smi
Tue Jul 15 14:27:47 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 546.30 Driver Version: 546.30 CUDA Version: 12.3 |
|-----------------------------------------+----------------------+----------------------+
| GPU Name TCC/WDDM | Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap | Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|=========================================+======================+======================|
| 0 NVIDIA GeForce RTX 3050 ... WDDM | 00000000:01:00.0 On | N/A |
| N/A 38C P8 4W / 40W | 590MiB / 4096MiB | 0% Default |
| | | N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=======================================================================================|
| 0 N/A N/A 3492 C+G ...CBS_cw5n1h2txyewy\TextInputHost.exe N/A |
| 0 N/A N/A 5940 C+G ...m Files\Mozilla Firefox\firefox.exe N/A |
| 0 N/A N/A 6636 C+G ...1.0_x64__8wekyb3d8bbwe\Video.UI.exe N/A |
| 0 N/A N/A 7016 C+G ...t.LockApp_cw5n1h2txyewy\LockApp.exe N/A |
| 0 N/A N/A 8300 C+G ...2txyewy\StartMenuExperienceHost.exe N/A |
| 0 N/A N/A 10404 C+G C:\Windows\explorer.exe N/A |
| 0 N/A N/A 11220 C+G ...cw5n1h2txyewy\CrossDeviceResume.exe N/A |
| 0 N/A N/A 11708 C+G ...5.0_x64__w2gh52qy24etm\Nahimic3.exe N/A |
| 0 N/A N/A 12048 C+G ...on\138.0.3351.77\msedgewebview2.exe N/A |
| 0 N/A N/A 12196 C+G ...nt.CBS_cw5n1h2txyewy\SearchHost.exe N/A |
| 0 N/A N/A 13148 C+G ...on\138.0.3351.77\msedgewebview2.exe N/A |
| 0 N/A N/A 14204 C+G ...5n1h2txyewy\ShellExperienceHost.exe N/A |
| 0 N/A N/A 15284 C+G ...al\Discord\app-1.0.9198\Discord.exe N/A |
| 0 N/A N/A 15740 C+G ...m Files\Mozilla Firefox\firefox.exe N/A |
| 0 N/A N/A 18280 C+G ...oogle\Chrome\Application\chrome.exe N/A |
| 0 N/A N/A 18600 C+G ...__8wekyb3d8bbwe\Notepad\Notepad.exe N/A |
| 0 N/A N/A 21080 C+G ...oogle\Chrome\Application\chrome.exe N/A |
| 0 N/A N/A 24208 C+G ...\cef\cef.win7x64\steamwebhelper.exe N/A |
| 0 N/A N/A 24260 C+G ...crosoft\Edge\Application\msedge.exe N/A |
| 0 N/A N/A 27912 C+G ...Programs\Microsoft VS Code\Code.exe N/A |
| 0 N/A N/A 30076 C+G ...273.0_x64__dt26b99r8h8gj\RtkUWP.exe N/A |
| 0 N/A N/A 34048 C+G ...s\System32\ApplicationFrameHost.exe N/A |
| 0 N/A N/A 43348 C+G C:\Windows\System32\ShellHost.exe N/A |
+---------------------------------------------------------------------------------------+

[jakevdp]

return [x.at](http://x.at/)[:].set(3.0)

Something is strange with your code example. Can you double-check that it runs as written?

[Cattaneo123]

Oh dear, it seems to have copied a little bit of the URL from my email, how odd.

from petsc4py import PETSc as PETSc
import jax
from functools import partial
import jax.numpy as jnp

jax.config.update("jax_enable_x64", True) 

@partial(jax.jit, donate_argnums=(0,))
def set_in_place(x):
    return x.at[:].set(3.0)

print('\nTesting jax from_dlpack given a PETSc vector that was allocated by PETSc')

x = jnp.ones((1000,1))
y_petsc = PETSc.Vec().createSeq(x.shape[0])
y_petsc.set(0.0)

print(hex(y_petsc.handle))

y2_petsc = PETSc.Vec().createWithDLPack(y_petsc.toDLPack('rw'))
y2_petsc.set(-1.0)

assert y_petsc.getValue(0) == y2_petsc.getValue(0)
print('After creating a second PETSc vector via a DLPack of the first, modifying the memory of one affects the other.')

y = jnp.from_dlpack(y_petsc, copy=False)

orig_ptr = y.unsafe_buffer_pointer()
print(f'before: ptr at {hex(orig_ptr)}')

y = set_in_place(y)

print(f'after:  ptr at {hex(y.unsafe_buffer_pointer())}')
assert orig_ptr == y.unsafe_buffer_pointer()

assert y_petsc.getValue(0) == y[0], f'The PETSc value {y_petsc.getValue(0)} did not match the JAX value {y[0]}, so modifying the JAX memory did not affect the PETSc memory.'

[jakevdp]

What is the output you're getting from running your script? (I'm not familiar with petsc4py – I tried following the install instructions but it hung on Getting requirements to build wheel ... and I gave up after ten minutes).

[jakevdp]

Also, since this concerns the copy semantics of JAX's dlpack support, #29963 may be relevant.

[Cattaneo123]

What is the output you're getting from running your script? (I'm not familiar with petsc4py – I tried following the install instructions but it hung on Getting requirements to build wheel ... and I gave up after ten minutes).

Be default, when trying to import using DLPack JAX says the memory is improperly aligned and a copy must be made. As such the pointers differ and the first assert fails. If I recompile PETSc to use a different alignment, such as 32 or 64 that error no longer appears and the first assert passes, so presumably JAX is successfully avoiding copies, however the second assert fails which means despite both the PETSc object and the JAX object (seemingly) referring to the same piece of memory the PETSc object isn't reflecting those changes. There's no error message or anything, it just doesn't seem to work.

Is there some limitation which prevents JAX methods from changing data that is owned by another package? What is odd is that the opposite, changing a JAX created vector that has been imported via DLPack to PETSc using PETSc functions, does reflect the changes in the JAX object.

As for the wheel building, that is weird as I've never seen it take that long. Did you build from source or use pip? Unless you already have PETSc installed for some reason I imagine pip which, if you have a local installation of PETSc as well, can sometimes misbehave.

[jakevdp]

There's no error message or anything, it just doesn't seem to work.

This may be the issue that #29963 addresses.

does reflect the changes in the JAX object.

You need to be really careful about this kind of thing. JAX is built on the fundamental assumption that buffers are immutable. When you work around this by exporting via dlpack, you break JAX's assumptions and may end up with wrong outputs.

Did you build from source or use pip?

I ran pip install petsc4py, as it recommends on the installation page. Still hung at 28 minutes

[Cattaneo123]

Ideally no changes to the data should be made outside of JAX functions, PETSc is only being used to pass around the data to different JAX routines, so there shouldn't be any changes outside of scope, I only bring it up as an example of ownership being passed around. Is it possible though that because the data is allocated by another program JAX will not allow a JITed function to perform in-place modification? If not then my next assumption has to be that PETSc disallows its objects from being modified outside of their scope.

Otherwise I'll see if I can find an error or warning in the underlying C/C++ section since I (foolishly) hadn't thought to check there.

[jakevdp]

If JAX does not donate the buffer, then assert orig_ptr == y.unsafe_buffer_pointer() will fail. It's still not clear to me whether it's this check or the next that is failing when you run the script.

[Cattaneo123]

Sorry, please forgive my lack of clarity. The second check (does the JAX view match the PETSc view) will always fail no matter what.

I can get the first assert to pass by changing the alignment of PETSc objects, but doing this does not fix the problem that it seems to not synch up the changes to the PETSc view, which makes me think even though it's not throwing an error regarding alignment it's still silently making a copy somewhere even though the alignment is the same and the PETSc pointer doesn't change.

My concern is that even though there are no errors or warning that say there is a copy being made, it seems a copy is being made. Because otherwise the change would be reflected in both views.

[jakevdp]

OK, thanks for the clarification. I think this is probably the same issue addressed by #29963 then: the issue fixed there is that jnp.from_dlpack is not properly respecting the value of the copy parameter, and so may be silently copying when you tell it not to.

[Cattaneo123]

OK, thanks for the clarification. I think this is probably the same issue addressed by #29963 then: the issue fixed there is that jnp.from_dlpack is not properly respecting the value of the copy parameter, and so may be silently copying when you tell it not to.

I see, thank you. Since it looks like it hasn't been merged I'll keep an eye on this thread and try again when/if it gets merged. Depending on what happens I'll try some other things.

Thank you very much for your time, you've been a big help.