feat: raw pointer capsule instantiation method

[lmmx]

• via [Idea] PyCapsule instantiators for raw pointer capsules #5688

Implements PyCapsule::new_pointer as described in the issue above

Purpose

Allows users to write much simpler code, without getting under the hood of pyo3-ffi to make a pyo3::PyCapsule

#[pyfunction(name = "rms_norm")]
fn rms_norm_jax(py: Python<'_>) -> PyResult<Bound<'_, PyCapsule>> {
    unsafe {
        PyCapsule::new_pointer(py, ffi::RmsNorm as *mut c_void, None)
    }
}

instead of

#[pyfunction(name = "rms_norm")]
fn rms_norm_jax(py: Python<'_>) -> PyResult<Bound<'_, PyCapsule>> {
    let fn_ptr: *mut c_void = ffi::RmsNorm as *mut c_void;
    let name = std::ptr::null();

    unsafe {
        let capsule = pyo3::ffi::PyCapsule_New(fn_ptr, name, None);
        if capsule.is_null() {
            return Err(pyo3::exceptions::PyRuntimeError::new_err(
                "Failed to create PyCapsule",
            ));
        }
        let any: Bound<'_, PyAny> = Bound::from_owned_ptr(py, capsule);
        Ok(any.cast_into_unchecked::<PyCapsule>())
    }
}

• Specifically this was the refactor enabled by switching to use this PR branch in jeertmans/extending-jax@8fbb5bb

[jeertmans]

I don't think you need to say that the pointer must be non-null, as this is checked internally by the function and will raise an error if the pointer NULL.

[davidhewitt]

Thanks for working on this, I think the proposed functionality makes sense. Some initial comments on the design below.

[davidhewitt]

How about using NonNull<c_void> here to make it clear in the API contract that the pointer must be non-null? Can then avoid checking for this in the code below, and defer to the caller to decide how they construct the NonNull.

[davidhewitt]

I suggest rather than leaking we make the name argument Option<&'static CStr>, given this is now trivial to construct.

I might even go further and recommend we always name capsules (this seems to be good practice), i.e. we could perhaps just have name: &'static CStr?

[davidhewitt]

It might be helpful to users to document that the destructor must be a raw function unlike the new_with_destructor case because there is nowhere to store a destructor closure in this form.

[davidhewitt]

Name bikeshedding, what do you think of new_with_pointer and new_with_pointer_and_destructor?

I somehow feel like new_pointer is describing that a "capsule pointer" gets returned rather than the capsule is constructed with a raw pointer.