Unable to register a customized GPU operator with CuPy

[leonlee-1]

Hi BrainPy Team,

A customized GPU operator of CuPy must be registered by calling bm.XLACustomOp(gpu_kernel=kernel) as described in the documentation (https://brainpy.readthedocs.io/tutorial_advanced/operator_custom_with_cupy.html), but this line of code causes AttributeError: module 'brainpy.math' has no attribute 'XLACustomOp'.

I've tried installing brainpy and also BrainX with cuda12 and cuda13 by following the instructions such as "pip install brainpy[cuda12]", but nothing worked. It seems that the cause is XLACustomOp doesn't exist any more. Please kindly shed some light on the resolutions.
Thank you.

[github-actions]

👋 Welcome to BrainPy! Thank you for opening your first issue.

We appreciate you taking the time to contribute to the project. A maintainer will review your issue and respond as soon as possible.

In the meantime, please make sure you've provided all the necessary information to help us understand and address your issue effectively.

If you're interested in contributing code, check out our Contributing Guide for more information on how to get started.

Happy coding! 🧠✨

[Routhleck]

Hi @leonlee-1,

You're correct — bm.XLACustomOp has been removed in recent versions of BrainPy. The functionality has been migrated to the separate brainevent package.

To register custom GPU operators with CuPy (or CPU operators with Numba), you should now use brainevent.XLACustomKernel directly. Here's how:

from brainevent import XLACustomKernel, warp_kernel

# Define your GPU kernel using Warp
kernel = warp_kernel(...)
custom_op = XLACustomKernel(
    name='my_custom_op',
    gpu_kernel=kernel
)

# Use it
result = custom_op(inputs, outs=[jax.ShapeDtypeStruct(shape, dtype)])

For detailed usage with CuPy/JIT, please refer to the brainevent documentation:

• brainevent.XLACustomKernel for general custom operators
• brainevent.numba_kernel for Numba CPU kernels
• brainevent.warp_kernel for Warp GPU kernels
• brainevent.pallas_kernel for Pallas GPU kernels

The old bm.XLACustomOp API was deprecated and removed during the BrainPy 2.7+ releases when brainevent was integrated as a separate dependency.

If you need help migrating specific code from the old API, feel free to ask!

[leonlee-1]

Hi @Routhleck,

I've tried creating a warp function and it works. Before that, I've spent several hours struggling with bm.XLACustomOp and reinstallation with various brainpy versions and so on. Thank you very much for your help.
Have a nice weekend.

[Routhleck]

Thank you for the confirmation, @leonlee-1! We're sorry for the confusion caused by the outdated documentation. We'll make sure to update the docs to reflect the migration to brainevent. Glad to hear the Warp approach is working for you! Have a great weekend as well!