Regarding package status

[SNMS95]

Hey @shwina

Thanks for this nice addition to the modern scientific computing.
We are building a simulation package in JAX and wanted to integrate into it pyamgx since having multi-grid solvers on GPU is a great addition.

However, it would be nice to know what the plans are for the repo?
i.e. Will it be actively maintained?
and is there a plan for a version release on pip ?

Regards

[shwina]

Hi @SNMS95 - thank you for reaching out. As you can tell, I'm not actively contributing to or adding new features to pyamgx, however--

• PRs are always welcome and I'll review/accept them
• I can make a best effort to resolve bugs that don't have a workaround
• No pip or conda packages are planned

For full disclosure, pyamgx has been a personal project of mine, and is not developed by the team working on the AMGX library at NVIDIA.

[vboussange]

Hey @SNMS95, we (@bhorowitz, @gdalle, myself...) are also quite interested in porting AMGX to JAX. Here are some leads:

1. Using the ffi feature of JAX to rewrite a JAX-specific wrapper. In that case, I see the advantage of developing a distinct library. I am not sure to have the capacity of leading that, but can definitely help.
2. Writing a method ptr_from_dlpack_interface similar to

   pyamgx/pyamgx/utils.pyx

   Lines 5 to 21 in 6229ff0

   	cdef uintptr_t ptr_from_array_interface(data, check_for_dtype=None) except *:
   	if hasattr(data, "__array_interface__"):
   	desc = data.__array_interface__
   	elif hasattr(data, "__cuda_array_interface__"):
   	desc = data.__cuda_array_interface__
   	else:
   	raise TypeError("Must pass array-like for data")
   	cdef uintptr_t ptr = desc["data"][0]
   	if check_for_dtype:
   	dtype = np.dtype(desc["typestr"])
   	if dtype != check_for_dtype:
   	raise ValueError("Invalid dtype: {}".format(dtype))
   	return ptr

   but leveraging the dlpack interface (see Support __cuda_array_interface__ on GPU  jax-ml/jax#1100 (comment)) to allow for zero-copy array transfer, and wrap pyamgx's solver.solve and relatives with a jax.pure_callback
3. Maybe that the __cuda_array_interface__ could work as is, see discussion here WIP: FFI support NIFTy-PPL/JAXbind#43

[SNMS95]

Hey @SNMS95, we (@bhorowitz, @gdalle, myself...) are also quite interested in porting AMGX to JAX. Here are some leads:

1. Using the ffi feature of JAX to rewrite a JAX-specific wrapper. In that case, I see the advantage of developing a distinct library. I am not sure to have the capacity of leading that, but can definitely help.

2. Writing a method ptr_from_dlpack_interface similar to

     [pyamgx/pyamgx/utils.pyx](https://github.com/shwina/pyamgx/blob/6229ff008ee5a264cfc1799eeb2f83d96da0aadc/pyamgx/utils.pyx#L5-L21)
   
   
       Lines 5 to 21
     in
     [6229ff0](/shwina/pyamgx/commit/6229ff008ee5a264cfc1799eeb2f83d96da0aadc)
   
   
   
   
   
       
         
          cdef uintptr_t ptr_from_array_interface(data, check_for_dtype=None) except *: 
       
   
       
         
              if hasattr(data, "__array_interface__"): 
       
   
       
         
                  desc = data.__array_interface__ 
       
   
       
         
              elif hasattr(data, "__cuda_array_interface__"): 
       
   
       
         
                  desc = data.__cuda_array_interface__ 
       
   
       
         
              else: 
       
   
       
         
                  raise TypeError("Must pass array-like for data") 
       
   
       
         
           
       
   
       
         
              cdef uintptr_t ptr = desc["data"][0] 
       
   
       
         
           
       
   
       
         
              if check_for_dtype: 
       
   
       
         
                  dtype = np.dtype(desc["typestr"]) 
       
   
       
         
           
       
   
       
         
                  if dtype != check_for_dtype: 
       
   
       
         
                      raise ValueError("Invalid dtype: {}".format(dtype)) 
       
   
       
         
           
       
   
       
         
              return ptr 
   

   but leveraging the dlpack interface (see Support cuda_array_interface on GPU  jax-ml/jax#1100 (comment)) to allow for zero-copy array transfer, and wrap pyamgx's solver.solve and relatives with a jax.pure_callback

3. Maybe that the __cuda_array_interface__ could work as is, see discussion here WIP: FFI support NIFTy-PPL/JAXbind#43

Hey guys,

totally missed this somehow.
I am not that well versed in the FFI aspect of JAX but
might be interesting for you? @johnviljoen

[johnviljoen]

Hi all! (And @gdalle again haha) I have done the work you need to do to instead port cuDSS to jax rather than amgx, but the requirements are the same from what I can tell.

I will be making that repo public in about a month when I manage to put something on arxiv and then it will be a template for how to do this (zero copy GPU arrays, jit, vmap, jvp). If that's too long, I'm happy to chat to talk about some of the implementation details/collaborate on getting this done for amgx!

[aziriahi]

@johnviljoen Nice meeting you here! I'm curious—what motivated you to port cuDSS to JAX? What application or problem are you trying to solve?

[johnviljoen]

@johannahaffner and I are working on getting a sparse clone of ipopt into JAX. We are fairly close to finishing - probably about a month to go, as a little taster we:

• generate sparsity patterns of JAX expressions with this
• solve the linear systems using cudss wrapped into JAX (will release with full work, and I aim to get it integrated into lineax if all goes well)
• will integrate the solver algorithm itself into a future version of optimistix
• we do not currently have matrix colouring implemented (relying on jax jit compiler for sparse backprop optimal choices), but we will probably implement something like @gdalle s sparsediffax for that