Includes basic np dlpack function that is tested to work across jax a… by apasarkar · Pull Request #1067 · fastplotlib/fastplotlib

apasarkar · 2026-07-30T21:42:52Z

The numpy dlpack fully sidesteps the need to rely on torch, but it requires numpy > 2.1. I think this is ok, since numpy is now at 2.5, and only users who care about high-performance (i.e. computing on GPUs) will use this execution path anyways.

…nd torch

apasarkar · 2026-07-31T03:33:22Z

@kushalkolar so one issue I've been seeing is that the cuda array interface execution path seems to produce spatially transposed data. See below for a reproducible example.


## Make a reproducible example:
data = np.random.rand(100, 50, 200)
data_torch = torch.as_tensor(data)

ref_ranges = {'time': (0, data.shape[0], 1)}
subplot_names = ['data_numpy', 'data_torch']
curr_ndw = fpl.NDWidget(ref_ranges,
                        names=subplot_names,
                        shape=(1, 2),
                        size=(800, 800))

dims = ('time', 'height', 'width')
spatial_dims = ('height', 'width')

ndw_vid_numpy = curr_ndw['data_numpy'].add_nd_image(data,
                                       dims,
                                       spatial_dims,
                                       name = 'data_numpy')

ndw_vid_torch = curr_ndw['data_torch'].add_nd_image(data_torch,
                                       dims,
                                       spatial_dims,
                                       name = 'data_torch')



curr_ndw.show()

The reason for this is because at the following line, the transpose does different things if the input is a tensor vs. if it is a numpy array.

fastplotlib/fastplotlib/widgets/nd_widget/_base.py

Line 537 in 91faf5b

return windowed_slice.transpose(*spatial_dims_int)

For a tensor, my_tensor.transpose(0, 1) will swap the axes 0 and 1. Otoh, my_numpy_array.transpose(0, 1) will keep the desired shape.

Is there any reason we are doing this final transpose step? For an individual ndprocessor, it does not seem necessary. I think there are two options:
(1) If there's no reason, drop it.
(2) If we need to do it, defer this to after we convert the dlpack-compatible array to numpy. This way the function call will be consistent.

kushalkolar

Where does it say that the from_dlpack requires numpy >= 2.1? Anyways that's fine, v2.1 is 2 years old at this point. pygfx requires min 2.1 anyways too https://github.com/pygfx/pygfx/blob/main/pyproject.toml#L28C1-L28C55

kushalkolar · 2026-07-31T06:58:19Z

-
-    return cupy.asnumpy(arr)
+
+    data = np.from_dlpack(arr, device='cpu') #This requires numpy >= 2.1


Suggested change

data = np.from_dlpack(arr, device='cpu') #This requires numpy >= 2.1

data = np.from_dlpack(arr, device='cpu')

kushalkolar · 2026-07-31T07:03:37Z

For a tensor, my_tensor.transpose(0, 1) will swap the axes 0 and 1. Otoh, my_numpy_array.transpose(0, 1) will keep the desired shape.

Is there any reason we are doing this final transpose step? For an individual ndprocessor, it does not seem necessary. I think there are two options: (1) If there's no reason, drop it. (2) If we need to do it, defer this to after we convert the dlpack-compatible array to numpy. This way the function call will be consistent.

Annoying that torch behavior is different from numpy.

The transpose is necessary by definition, that's how the NDWidget spec is. The specified spatial dims order defines the dim order for the graphic data.

Doing the transpose after it's converted to numpy sounds like the way to go.

apasarkar · 2026-08-02T02:53:22Z

@kushalkolar Updated this -- so I've delayed the permutation operation until the sliced data is guaranteed to be a numpy array. nd vectors processor, nd image processor, and nd positions processor all subclass ndprocessor, so those are the files where i made this change.

One thing to note: with this setup, the spatial functions must be written so they are applied data that has not yet been permuted. for example, if i specify data like this:

dims = ['time', 'a', 'b']
spatial_dims = ['b', 'a'] # This means we end up transposing the sliced data

The spatial window function received the spatial dimensions in order (a, b) with these changes.

I actually think this is ok -- this forces the user to write all their window/spatial function logic using the dimension order specified by dims, but the final order is transposed based on what's in spatial_dims. So we would ofc need to update the specs. I think so long as the user can write their whole pipeline (raw data --> displayed graphic) with a fixed dimension ordering it will be clear.

Let me know what you think!

kushalkolar · 2026-08-04T19:22:19Z

One thing to note: with this setup, the spatial functions must be written so they are applied data that has not yet been permuted. for example, if i specify data like this:

dims = ['time', 'a', 'b'] spatial_dims = ['b', 'a'] # This means we end up transposing the sliced data

The spatial window function received the spatial dimensions in order (a, b) with these changes.

I actually think this is ok -- this forces the user to write all their window/spatial function logic using the dimension order specified by dims, but the final order is transposed based on what's in spatial_dims. So we would ofc need to update the specs. I think so long as the user can write their whole pipeline (raw data --> displayed graphic) with a fixed dimension ordering it will be clear.

Let me know what you think!

I think that makes sense. The whole point is that window_funcs and spatial_funcs should be fast and async, and things are naturally async when they remain in torch as long as possible. The spatial_dims arg for an NDGraphic is really just supposed to be the display order of the dimensions for rendering the graphic, the [x, y, z] dims for positional data, [row, col] for image data nad [depth, row, col] for volumetric image data. The user needs to know what their dims are to create their window_funcs -> spatial_funcs pipeline.

Maybe we should rename the arg to display_order or display_dims or render_dims or something?

apasarkar · 2026-08-04T19:39:04Z

ordered_spatial_dims might be a good one, specifying the two roles of this parameter: (1) it specifies what the spatial dims are and (2) it gives an order to output the spatial dims. thoughts?

apasarkar · 2026-08-04T19:41:42Z

@kushalkolar I'll also flag the following potential issue:

fastplotlib/fastplotlib/widgets/nd_widget/_base.py

Line 525 in 91faf5b

windowed_slice = windowed_slice.squeeze(axis=slider_dims_int)

At this line, a function "squeeze" is used. On a first read, I expected it to fail with torch tensors, since the documented parameters are tensor.squeeze(dim = ...), whereas the code is doing .squeeze(axis =...). Turns out at some point in pytorch the developers allowed axis to make it interoperable with numpy.

Anyways it's not clear with other libraries that are compatible with dlpack whether this functionality will break.

kushalkolar · 2026-08-04T20:03:59Z

cupy is supposed to be just like numpy and accepts axis: https://docs.cupy.dev/en/stable/reference/generated/cupy.squeeze.html
jax likewise: https://docs.jax.dev/en/latest/_autosummary/jax.numpy.squeeze.html

where is it documented that torch also accepts axis? I don't see it here: https://docs.pytorch.org/docs/2.13/generated/torch.squeeze.html#torch.squeeze

apasarkar · 2026-08-05T02:16:31Z

@kushalkolar Example:

https://discuss.pytorch.org/t/confuse-about-how-torch-argmax-supporting-axis-parameter-inplace-of-dim/182592

Looks like under the hood there is some code that parses keywords.

kushalkolar · 2026-08-05T02:19:13Z

@kushalkolar Example:

https://discuss.pytorch.org/t/confuse-about-how-torch-argmax-supporting-axis-parameter-inplace-of-dim/182592

Looks like under the hood there is some code that parses keywords.

wow that's hidden deep

apasarkar · 2026-08-05T15:04:21Z

@kushalkolar I'm thinking from the standpoint of documentation. We ideally want to be able to tell someone "your dlpack compliant array can work here". Maybe there is not a workaround with the axis squeeze thing though

kushalkolar · 2026-08-05T15:09:51Z

If it works with torch, cupy and Jax that covers most use cases I'd think?

apasarkar · 2026-08-05T16:17:14Z

@kushalkolar Fair enough, we can just make sure to note this in the documentation.

Ok so what's left here - just the rename of the spatial_dims param to something else? If you're ok with renaming is ordered_spatial_dims, we'll have to update many of the function signatures in ndwidget, so let's make sure we're both happy with the naming here and I can go ahead and do it.

kushalkolar · 2026-08-05T16:32:29Z

I added it to the list of renames to do later in bulk on #971

kushalkolar · 2026-08-05T16:33:32Z

what's left is the suggestion and updating pyproject.toml and merge latest ndwidget into this

apasarkar · 2026-08-05T17:47:34Z

@kushalkolar Added the suggested change.

Re: pyproject.toml, we are now using np.from_dlpack exclusively to bring gpu arrays to cpu and from there do zero-cost conversion to numpy. So it looks like we don't need to change the pyproject.toml at all.

kushalkolar · 2026-08-05T17:54:49Z

@kushalkolar Added the suggested change.

Re: pyproject.toml, we are now using np.from_dlpack exclusively to bring gpu arrays to cpu and from there do zero-cost conversion to numpy. So it looks like we don't need to change the pyproject.toml at all.

For the min version pin?

…sion across devices

apasarkar · 2026-08-05T18:08:14Z

@kushalkolar sorry I thought I did that earlier haha.
I think g2g in latest commit.

kushalkolar · 2026-08-05T18:19:29Z

+        # Axis order of the spatial dimensions to display
+        self._spatial_dims_int = tuple(
+            self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims
+        )
+


spatial_dims is a mutable property, so this should be a read only property or set in the spatial_dims setter as a private attribute (probably preferable, it doesn't have to be public) and not created once in the constructor

@kushalkolar good point, updated in latest commit. I also renamed spatial_dims_int to spatial_dims_indices for clarity

… changed

kushalkolar · 2026-08-06T05:20:45Z

+        """
+        The ordered sequence of data indices that will be displayed
+        """
+        return self._spatial_dims_indices


Suggested change

return self._spatial_dims_indices

return tuple(

self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims

)

kushalkolar · 2026-08-06T05:21:13Z

+    @property
+    def spatial_dims_indices(self) -> tuple[int, ...]:
+        """
+        The ordered sequence of data indices that will be displayed


Suggested change

The ordered sequence of data indices that will be displayed

ordered spatial dim indices that correspond to the named spatial dims

kushalkolar · 2026-08-06T05:21:34Z

+        ## This is the ordered sequence of data indices that will be displayed
+        self._spatial_dims_indices = tuple(
+            self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims
+        )
+


Suggested change

## This is the ordered sequence of data indices that will be displayed

self._spatial_dims_indices = tuple(

self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims

)

kushalkolar · 2026-08-06T05:24:59Z

+        ## This is the ordered sequence of data indices that will be displayed
+        self._spatial_dims_indices = tuple(
+            self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims
+        )
+
+    @property
+    def spatial_dims_indices(self) -> tuple[int, ...]:
+        """
+        The ordered sequence of data indices that will be displayed
+        """
+        return self._spatial_dims_indices
+


Suggested change

## This is the ordered sequence of data indices that will be displayed

self._spatial_dims_indices = tuple(

self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims

)

@property

def spatial_dims_indices(self) -> tuple[int, ...]:

"""

The ordered sequence of data indices that will be displayed

"""

return self._spatial_dims_indices

kushalkolar · 2026-08-06T05:25:04Z

            window_output = await run_in_thread_pool(self._executor, cuda_to_numpy, window_output)

-        return window_output
+


Suggested change

kushalkolar · 2026-08-06T05:25:13Z

+        ## This is the ordered sequence of data indices that will be displayed
+        self._spatial_dims_indices = tuple(
+            self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims
+        )
+
+    @property
+    def spatial_dims_indices(self) -> tuple[int, ...]:
+        """
+        The ordered sequence of data indices that will be displayed
+        """
+        return self._spatial_dims_indices
+


Suggested change

## This is the ordered sequence of data indices that will be displayed

self._spatial_dims_indices = tuple(

self.spatial_dims.index(d) for d in self.dims if d in self.spatial_dims

)

@property

def spatial_dims_indices(self) -> tuple[int, ...]:

"""

The ordered sequence of data indices that will be displayed

"""

return self._spatial_dims_indices

kushalkolar · 2026-08-06T05:25:35Z

            window_output = await run_in_thread_pool(self._executor, cuda_to_numpy, window_output)

-        return window_output
+


Suggested change

kushalkolar · 2026-08-06T05:26:01Z


-        return window_output
+
+        return window_output.transpose(*self._spatial_dims_indices)


Suggested change

return window_output.transpose(*self._spatial_dims_indices)

return window_output.transpose(*self.spatial_dims_indices)

it's a public property use it

…in the base class, updates docs, uses public property in the ndprocessor subclasses

apasarkar · 2026-08-06T15:41:09Z

@kushalkolar thanks -- all makes sense. made the changes, spatial_dims_indices is computed and set in the public property of the base class now. i was hesitant to do it this way at the start because the computation to define spatial dims indices will now be run at every single getitem. but if the number of dimensions is small, the overhead should never be a problem.

kushalkolar · 2026-08-06T17:51:26Z

It's probably a nanosecond scale operation.

Includes basic np dlpack function that is tested to work across jax a…

2f232d6

…nd torch

apasarkar requested review from clewis7 and kushalkolar as code owners July 30, 2026 21:42

apasarkar marked this pull request as draft July 31, 2026 02:58

kushalkolar reviewed Jul 31, 2026

View reviewed changes

delays dimension permutation for all code that uses ndprocessor

9621942

kushalkolar mentioned this pull request Aug 5, 2026

NDWidget #971

Draft

22 tasks

Eliminates old comment

14cd6e2

Updates the numpy dep to be the min version that allows dlpack conver…

cfee3a1

…sion across devices

kushalkolar requested changes Aug 5, 2026

View reviewed changes

Includes code to update spatial dims indices whenever spatial dims is…

a2d273a

… changed

kushalkolar requested changes Aug 6, 2026

View reviewed changes

Moves the spatial dims indices computation to the read only property …

a1512f6

…in the base class, updates docs, uses public property in the ndprocessor subclasses

kushalkolar marked this pull request as ready for review August 7, 2026 00:18

kushalkolar merged commit a2b9d93 into ndwidget Aug 7, 2026

kushalkolar deleted the drop_cupy branch August 7, 2026 00:19

kushalkolar mentioned this pull request Aug 7, 2026

use torch dlpack instead of cupy to convert cuda array to numpy #1063

Closed


		return cupy.asnumpy(arr)

		data = np.from_dlpack(arr, device='cpu') #This requires numpy >= 2.1

	The ordered sequence of data indices that will be displayed
	ordered spatial dim indices that correspond to the named spatial dims

		window_output = await run_in_thread_pool(self._executor, cuda_to_numpy, window_output)

		return window_output


		return window_output

		return window_output.transpose(*self._spatial_dims_indices)

	return window_output.transpose(*self._spatial_dims_indices)
	return window_output.transpose(*self.spatial_dims_indices)

Uh oh!

Conversation

apasarkar commented Jul 30, 2026

Uh oh!

apasarkar commented Jul 31, 2026

Uh oh!

kushalkolar left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kushalkolar commented Jul 31, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

apasarkar commented Aug 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kushalkolar commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

apasarkar commented Aug 4, 2026

Uh oh!

apasarkar commented Aug 4, 2026

Uh oh!

kushalkolar commented Aug 4, 2026

Uh oh!

apasarkar commented Aug 5, 2026

Uh oh!

kushalkolar commented Aug 5, 2026

Uh oh!

apasarkar commented Aug 5, 2026

Uh oh!

kushalkolar commented Aug 5, 2026

Uh oh!

apasarkar commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

kushalkolar commented Aug 5, 2026

Uh oh!

kushalkolar commented Aug 5, 2026

Uh oh!

apasarkar commented Aug 5, 2026

Uh oh!

kushalkolar commented Aug 5, 2026

Uh oh!

apasarkar commented Aug 5, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

apasarkar commented Aug 6, 2026

Uh oh!

kushalkolar commented Aug 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

kushalkolar commented Jul 31, 2026 •

edited

Loading

apasarkar commented Aug 2, 2026 •

edited

Loading

kushalkolar commented Aug 4, 2026 •

edited

Loading

apasarkar commented Aug 5, 2026 •

edited

Loading