<html>
<head>
<meta http-equiv="Content-Type" content="text/html; charset=utf-8">
<style type="text/css" style="display:none;"> P {margin-top:0;margin-bottom:0;} </style>
</head>
<body dir="ltr">
<div class="elementToProof" style="font-family: Aptos, Aptos_EmbeddedFont, Aptos_MSFontService, Calibri, Helvetica, sans-serif; font-size: 12pt; color: rgb(0, 0, 0);">
It is expected that changes made in Jax are not reflected in the PETSc object. The issue has been explained in my previous message (point 2).</div>
<div style="font-family: Aptos, Aptos_EmbeddedFont, Aptos_MSFontService, Calibri, Helvetica, sans-serif; font-size: 12pt; color: rgb(0, 0, 0);">
<br>
</div>
<div style="font-family: Aptos, Aptos_EmbeddedFont, Aptos_MSFontService, Calibri, Helvetica, sans-serif; font-size: 12pt; color: rgb(0, 0, 0);">
Hong</div>
<div id="appendonsend"></div>
<hr style="display:inline-block;width:98%" tabindex="-1">
<div id="divRplyFwdMsg" dir="ltr"><font face="Calibri, sans-serif" style="font-size:11pt" color="#000000"><b>From:</b> Alberto Cattaneo <bubu.cattaneo@gmail.com><br>
<b>Sent:</b> Tuesday, July 15, 2025 1:13 PM<br>
<b>To:</b> Zhang, Hong <hongzhang@anl.gov><br>
<b>Subject:</b> Re: [petsc-users] Petsc/Jax no copy interfacing issues</font>
<div> </div>
</div>
<style>
<!--
#x_pfptBannerb00ozls
{display:block!important;
visibility:visible!important;
opacity:1!important;
background-color:#D0D8DC!important;
max-width:none!important;
max-height:none!important}
-->
</style>
<div>
<div style="display:none!important; display:none; visibility:hidden; font-size:1px; color:#ffffff; line-height:1px; height:0px; max-height:0px; opacity:0; overflow:hidden">
Odd, I was using double precision (forgot to include that in the example, sorry) but on my machineI’m still not seeing the changes made reflected in the PETSc object. Are the changes reflected on your end? Is it possibly an ownership issue?</div>
<div style="display:none!important; display:none; visibility:hidden; font-size:1px; color:#ffffff; line-height:1px; max-height:0px; opacity:0; overflow:hidden">
ZjQcmQRYFpfptBannerStart</div>
<div dir="ltr" id="x_pfptBannerb00ozls" style="display:block!important; text-align:left!important; margin:16px 0px 16px 0px!important; padding:8px 16px 8px 16px!important; border-radius:4px!important; min-width:200px!important; background-color:#D0D8DC!important; background-color:#D0D8DC; border-top:4px solid #90a4ae!important; border-top:4px solid #90a4ae">
<div id="x_pfptBannerb00ozls" style="float:left!important; display:block!important; margin:0px 0px 1px 0px!important; max-width:600px!important">
<div id="x_pfptBannerb00ozls" style="display:block!important; visibility:visible!important; background-color:#D0D8DC!important; color:#000000!important; color:#000000; font-family:'Arial',sans-serif!important; font-family:'Arial',sans-serif; font-weight:bold!important; font-weight:bold; font-size:14px!important; line-height:18px!important; line-height:18px">
This Message Is From an External Sender </div>
<div id="x_pfptBannerb00ozls" style="display:block!important; visibility:visible!important; background-color:#D0D8DC!important; color:#000000!important; color:#000000; font-weight:normal; font-family:'Arial',sans-serif!important; font-family:'Arial',sans-serif; font-size:12px!important; line-height:18px!important; line-height:18px; margin-top:2px!important">
This message came from outside your organization. </div>
</div>
<div style="clear:both!important; display:block!important; visibility:hidden!important; line-height:0!important; font-size:0.01px!important; height:0px">
</div>
</div>
<div style="display:none!important; display:none; visibility:hidden; font-size:1px; color:#ffffff; line-height:1px; max-height:0px; opacity:0; overflow:hidden">
ZjQcmQRYFpfptBannerEnd</div>
<div dir="ltr">Odd, I was using double precision (forgot to include that in the example, sorry) but on my machineI’m still not seeing the changes made reflected in the PETSc object. Are the changes reflected on your end? Is it possibly an ownership issue?</div>
<img alt="" width="0" height="0" style="width:2px; max-height:0; overflow:hidden" src="https://email-signature-image.com/signature.gif?u=140704&e=848292714&v=54d2610819e15f0b9417e49c9dc5640e57a6e827203f1f930877896fa6296df3"><br>
<div class="x_gmail_quote x_gmail_quote_container">
<div dir="ltr" class="x_gmail_attr">On Tue, Jul 8, 2025 at 3:56 PM Zhang, Hong <<a href="mailto:hongzhang@anl.gov">hongzhang@anl.gov</a>> wrote:<br>
</div>
<blockquote class="x_gmail_quote" style="margin:0px 0px 0px 0.8ex; border-left:1px solid rgb(204,204,204); padding-left:1ex">
<div class="x_msg-6707536908317760551">
<div lang="EN-US" style="">
<div class="x_m_-6707536908317760551WordSection1">
<p>Hi Alberto,<br>
<br>
1. To check the array pointer on the PETSc side, you can do print(hex(y_petsc.array.ctypes.data)). Then you will see a pointer mismatch caused by the line
<span style="color:rgb(49,49,49)">y = jnp.from_dlpack(y_petsc, copy=False). This is because you configured PETSc in double precision, but JAX uses single precision by default. You can either add</span>
<span style="color:rgb(49,49,49)">jax.config.update("jax_enable_x64", True) to make JAX use double precision number or configure PETSc to support single precision.
<br>
<br>
2. Once you fix this precision mismatch, the in-place conversion between PETSc and JAX should work. However, .at[].set() in JAX does not guarantee to operate in-place. The array updates in JAX are generally performed out-of-place by design. You may do the updates
in PETSc so that it won’t break the zero-copy system.</span><u></u><u></u></p>
<p class="x_MsoNormal">Hong<u></u><u></u></p>
<p class="x_MsoNormal"><u></u> <u></u></p>
<div style="border-right:none; border-bottom:none; border-left:none; border-top:1pt solid rgb(181,196,223); padding:3pt 0in 0in">
<p class="x_MsoNormal"><b><span style="font-family:Calibri,sans-serif; color:black">From:
</span></b><span style="font-family:Calibri,sans-serif; color:black">petsc-users <<a href="mailto:petsc-users-bounces@mcs.anl.gov" target="_blank">petsc-users-bounces@mcs.anl.gov</a>> on behalf of Alberto Cattaneo <<a href="mailto:bubu.cattaneo@gmail.com" target="_blank">bubu.cattaneo@gmail.com</a>><br>
<b>Date: </b>Monday, July 7, 2025 at 8:40 AM<br>
<b>To: </b>"<a href="mailto:petsc-users@mcs.anl.gov" target="_blank">petsc-users@mcs.anl.gov</a>" <<a href="mailto:petsc-users@mcs.anl.gov" target="_blank">petsc-users@mcs.anl.gov</a>><br>
<b>Subject: </b>[petsc-users] Petsc/Jax no copy interfacing issues<u></u><u></u></span></p>
</div>
<div>
<p class="x_MsoNormal"><u></u> <u></u></p>
</div>
<div>
<div>
<p class="x_MsoNormal">Greetings.<u></u><u></u></p>
<div>
<p class="x_MsoNormal">I hope this email reaches you well. I’m 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.<u></u><u></u></p>
</div>
<div>
<p class="x_MsoNormal"><u></u> <u></u></p>
</div>
<div>
<p class="x_MsoNormal">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. <u></u><u></u></p>
</div>
<div>
<p class="x_MsoNormal"><u></u> <u></u></p>
</div>
<div>
<p class="x_MsoNormal">I’ve included a code snippet below:<u></u><u></u></p>
</div>
<div>
<p class="x_MsoNormal"><u></u> <u></u></p>
</div>
<div>
<div>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">from petsc4py import PETSc as PETSc</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">import jax</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">from functools import partial</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">import jax.numpy as jnp</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="word-spacing:1px"><i><span style="color:rgb(49,49,49)"> </span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">@partial(jax.jit, donate_argnums=(0,))</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">def set_in_place(x):</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)"> return <a href="https://urldefense.us/v3/__https://gcc02.safelinks.protection.outlook.com/?url=https*3A*2F*2Furldefense.us*2Fv3*2F__http*3A*2Fx.at__*3B!!G_uCfscf7eWS!cqxG3TobpS7WZAgzxjrlWaxhAiiwWk4i9-WKReIWrc04LoXg4Y8zCkEDYGm_l5GilInGXbyzJWrD3BPRaTPlZHhIdz33*24&data=05*7C02*7Cpetsc-users*40mcs.anl.gov*7Cd16ee49fac3e47f97dc408ddc4b64c19*7C0cfca18525f749e38ae7704d5326e285*7C0*7C0*7C638883009407911569*7CUnknown*7CTWFpbGZsb3d8eyJFbXB0eU1hcGkiOnRydWUsIlYiOiIwLjAuMDAwMCIsIlAiOiJXaW4zMiIsIkFOIjoiTWFpbCIsIldUIjoyfQ*3D*3D*7C0*7C*7C*7C&sdata=d9kaMxlXfBWQZwzB4IVvYeLz7Ru2LI64jLkBWJZnGHM*3D&reserved=0__;JSUlJSUlJSUlJSUlJSUlJSUlJSUlJSUlJSU!!G_uCfscf7eWS!bFuLQebXmxQbeA2pyUDtq86jpMg21z70S5aqiiqlj5t4fnkwEZVop06kAVrYZpzLLQ82f61tj5fSFWXdKUWRvbcB9Q$" originalsrc="https://urldefense.us/v3/__http:/x.at__;!!G_uCfscf7eWS!cqxG3TobpS7WZAgzxjrlWaxhAiiwWk4i9-WKReIWrc04LoXg4Y8zCkEDYGm_l5GilInGXbyzJWrD3BPRaTPlZHhIdz33$" originalsrc="https://urldefense.us/v3/__http:/x.at__;!!G_uCfscf7eWS!cqxG3TobpS7WZAgzxjrlWaxhAiiwWk4i9-WKReIWrc04LoXg4Y8zCkEDYGm_l5GilInGXbyzJWrD3BPRaTPlZHhIdz33$" target="_blank"><span style="color:rgb(66,133,244)">x.at</span></a>[:].set(3.0)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="word-spacing:1px"><i><span style="color:rgb(49,49,49)"> </span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">print('\nTesting jax from_dlpack given a PETSc vector that was allocated by PETSc')</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">x = jnp.ones((1000,1))</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">y_petsc = PETSc.Vec().createSeq(x.shape[0])</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">y_petsc.set(0.0)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">print(hex(y_petsc.handle))</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">y2_petsc = PETSc.Vec().createWithDLPack(y_petsc.toDLPack('rw'))</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">y2_petsc.set(-1.0)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">assert y_petsc.getValue(0) == y2_petsc.getValue(0)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">print('After creating a second PETSc vector via a DLPack of the first, modifying the memory of one affects the other.')</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">#y = jnp.from_dlpack(y_petsc.toDLPack('rw'), copy=False)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">y = jnp.from_dlpack(y_petsc, copy=False)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">orig_ptr = y.unsafe_buffer_pointer()</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">print(f'before: ptr at {hex(orig_ptr)}')</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">y = set_in_place(y)</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">print(f'after: ptr at {hex(y.unsafe_buffer_pointer())}')</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">assert orig_ptr == y.unsafe_buffer_pointer()</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><i><span style="color:rgb(49,49,49)">#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.'</span></i><span style="color:rgb(49,49,49)"><u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><span style="color:rgb(49,49,49)"><u></u> <u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><span style="color:rgb(49,49,49)">I’d like the bottom two asserts to pass, but I can only get one of them. If somebody is familiar with this issue I’d greatly appreciate the assistance.<u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><span style="color:rgb(49,49,49)">Respectfully:<u></u><u></u></span></p>
<p style="font-size:1rem; word-spacing:1px"><span style="color:rgb(49,49,49)">Alberto<u></u><u></u></span></p>
</div>
</div>
</div>
</div>
</div>
</div>
</div>
</blockquote>
</div>
</div>
</body>
</html>