gridpoints 1.0.6__tar.gz → 1.1.0__tar.gz

gridpoints-1.1.0/PKG-INFO

@@ -0,0 +1,116 @@
+Metadata-Version: 2.4
+Name: gridpoints
+Version: 1.1.0
+Summary: fast multidimensional sorting tool for point clouds using the monotonic lagrangian grid 
+Author: Armand de Cacqueray
+License: MIT
+Project-URL: Homepage, https://neighborhood-grid.github.io/gridpoints/
+Project-URL: Repository, https://github.com/Neighborhood-Grid/gridpoints
+Keywords: sorting,point-clouds,multi-indexing,grid,neighborhood-grid,monotonic-lagrangian-grid,squarenet,nearest-neighbor
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Provides-Extra: demo
+Requires-Dist: squarenet; extra == "demo"
+Dynamic: license-file
+
+## ❒ CubeNet: Multi dimensional sort for point clouds 
+# with the gridpoints python package
+
+`Gridpoints` maps raw indexes of unstructured **point clouds** to structured grids multi-index  through a **bijective transformation**: 
+One point, one cell, 
+```text
+n <-> [i,j,…](n).
+```
+It replaces and enhances the [squarenet](https://github.com/Neighborhood-Grid/SquareNet) project with a more robust points sorting algorithm. The goal is 
+to build a spatialy coherent multi index to later store and query the point cloud efficiently, based on adaptative axes (lines, columns, and more in 3D+) as showed on following exemple :
+
+<img src="https://raw.githubusercontent.com/Neighborhood-Grid/CubeNet/main/ballexemple.png">
+
+What it does: Take raw point cloud `P(N, D)`  and find a grid shape and an index permutation `order` such that `Pgrid(I, J, …, D)` = P[order].reshape(*gridshape, D) is sorted along every axis of the grid. E.g. in 3D, for Pgrid = (x, y, z):
+```text
+x[i+1, j, k] >= x[i, j, k]
+y[i, j+1, k] >= y[i, j, k]
+z[i, j, k+1] >= z[i, j, k]
+```
+
+→ On the `Pgrid` view of `P`, neighbor queries become a simple stencil look-up
+```text
+neighborhood(Pgrid[i, j, k]) = {Pgrid[i±di, j±dj, k±dk] | (di, dj, dk) ≤ R},
+````
+where `R` is a radius cutoff to determine, allowing local operations in O(N) time.  
+→ Geometric operations (convolution, clustering, interpolation …) can then be applied **in O(N) time** directly on the `Pgrid` view instead of relying on complex graph or point-cloud methods
+
+`P` can be a NumPy, PyTorch or CuPy array of any dimension (N, D). To allow natural padding when the grid has more cells than they are points in the cloud, NaNs and Infs are supported in a consistent manner: 
+- nans → random position
+- (+-) infs → border of the grid
+  
+This allow to gridsort prime or variable number of points, as long as one is ready to deal with void/special grid cells.
+
+Expected runtime for sorting 1 million points, tested on google colab GPU T4 for 2D/3D point clouds: 
+- Numpy: 10s
+- Torch (gpu): 500 ms
+- Cupy (gpu): 1 s (after compilation e.g. not on the first cold call, which take 10 s)
+
+`When not to use gridpoints ?`
+- high dimension: the package is implemented to support arbitrary dimension, but sweetspot is really 2D/3D. dimensions 4-6 might still be reasonable depending on the task, but anything above 8D is generally too high dimensional for gridpoints.
+- weird geometries. Supported datasets goes beyond smooth convex distributions: map of Indonesia, a sponge, a donuts, an elephant, an eggshell (by specifying a suitable gridshape e.g. (128,128,2) to `gridpoints.sort()`). You can see various exemples in the [plot](https://github.com/Neighborhood-Grid/CubeNet/blob/main/plots) folder But with some limits. Bad fits: a spider web, a wind turbine, same eggshell with a naive 3d sort (gridshape = (32,32,32)) ... the issue is not that gridpoints cannot sort these distributions, but that it will produce a poor representation of the geometry.
+- small point clouds: beyond a few hundred points, local operations in O(N) time is not worth the overhead, because naive quadratic implementations will probably be simultaneously simpler and faster.
+
+---
+
+### Installation
+
+```bash
+pip install gridpoints          # core only
+pip install gridpoints[demo]    # for the demonstration notebook.
+```
+Full Demo: [notebook.ipynb](https://github.com/Neighborhood-Grid/CubeNet/blob/main/notebook.ipynb)
+
+### Quickstart
+
+```python
+import gridpoints as grid
+import numpy as np
+
+# Raw point cloud (numpy, pytorch or cupy)
+A = np.random.rand(1_000_000, 3)
+
+# Sorted view: place the points inside the grid
+order = grid.argsort(A, gridshape=(100, 100, 100))
+Bflat = A[order] #flat but sorted with a grid layout (C-order)
+Bgrid = Bflat.reshape(100, 100, 100, 3) #reshaped as a grid 
+
+# Rest of your pipeline, working with grids
+Cgrid = apply_something(Bgrid) #(100, 100, 100, *)
+
+# Back to the original points indexing
+Cflat = Cgrid.reshape(1_000_000, *)
+orderinv = grid.invert_permutation(order)
+C = Cflat[orderinv]   # matches the initial points order
+```
+
+---
+
+### Note on rectangular gridshapes
+
+The eggshell example illustrates the importance of selecting a suitable grid-shape decomposition. This can itself be challenging when the point cloud is assumed to follow a complex topological pattern, such as a thin surface rather than a full volumetric domain. While the precise decomposition is not particularly important (e.g. (18,20,16) vs. (16,15,24) will generally make little difference), the orders of magnitude of the different dimensions do matter and should be tuned to the dataset: (18,20,16) and (36,40,4) can lead to substantially different results. As mentioned, nan/inf padding can help accomodate integer-factorization constraints
+
+Another possibility for smoothing out topology-specific effects is to perform the sort after randomly projecting the dataset onto a lower-dimensional subspace. Random projections are known to approximately preserve pairwise distances, as formalized by the Johnson–Lindenstrauss lemma.
+
+
+### Note on the cutoff radius R
+
+There is no strict theoretical guarantee about what the cutof radius R should be for a given task. E.g the relative grid position between a point and its nearest neighbors can't be garanted to be in the exact adjacent grid cells. What is guaranteed from the sorted ordering is only **grid monotonicity**: *x* coordinates increase along rows, *y* coordinates along columns, and so on.
+
+As an example, empirical results in 2-D show that `R = 5` is enough for ~99 % of the nearest neighbors; some outlier neighbors will sit further apart for complex geometries with pronounced peaks, holes or any non-smoothness. When a stricter neighborhood is required, or in high dimensional setting, the solution might be to build an assembly of grid experts, each working on a rotated / projected view of the points, as discussed in [this topic](https://github.com/glotzerlab/freud/discussions/1417). An other possibility would be to tile the grid: each tile is enhanced with location metadata (e.g. a bounding box), allowing for pruning pairs of tiles with a distance certified to be far enough for a given criterion.
+
+
+### Note on efficient stencil operations
+
+The typical use-case of `Gridpoints` is to allow fast local operations on arbitrary point clouds using stencil kernels:
+
+```text
+output(i, j, k) = f( Pgrid[i±di, j±dj, k±dk] | di, dj, dk in local window )
+```
+To go beyond standard (slow) python loops, this kind of kernel computation can be accelerated with native grid convolution operations of standard libraries whenever possible, or with `pystencils` or `taichi` compilers for complex/non O(N) grid kernels. On GPU, tiling the grid and leveraging a custom `triton` kernel might be particularly efficient.

gridpoints-1.1.0/README.md

@@ -0,0 +1,100 @@
+## ❒ CubeNet: Multi dimensional sort for point clouds 
+# with the gridpoints python package
+
+`Gridpoints` maps raw indexes of unstructured **point clouds** to structured grids multi-index  through a **bijective transformation**: 
+One point, one cell, 
+```text
+n <-> [i,j,…](n).
+```
+It replaces and enhances the [squarenet](https://github.com/Neighborhood-Grid/SquareNet) project with a more robust points sorting algorithm. The goal is 
+to build a spatialy coherent multi index to later store and query the point cloud efficiently, based on adaptative axes (lines, columns, and more in 3D+) as showed on following exemple :
+
+<img src="https://raw.githubusercontent.com/Neighborhood-Grid/CubeNet/main/ballexemple.png">
+
+What it does: Take raw point cloud `P(N, D)`  and find a grid shape and an index permutation `order` such that `Pgrid(I, J, …, D)` = P[order].reshape(*gridshape, D) is sorted along every axis of the grid. E.g. in 3D, for Pgrid = (x, y, z):
+```text
+x[i+1, j, k] >= x[i, j, k]
+y[i, j+1, k] >= y[i, j, k]
+z[i, j, k+1] >= z[i, j, k]
+```
+
+→ On the `Pgrid` view of `P`, neighbor queries become a simple stencil look-up
+```text
+neighborhood(Pgrid[i, j, k]) = {Pgrid[i±di, j±dj, k±dk] | (di, dj, dk) ≤ R},
+````
+where `R` is a radius cutoff to determine, allowing local operations in O(N) time.  
+→ Geometric operations (convolution, clustering, interpolation …) can then be applied **in O(N) time** directly on the `Pgrid` view instead of relying on complex graph or point-cloud methods
+
+`P` can be a NumPy, PyTorch or CuPy array of any dimension (N, D). To allow natural padding when the grid has more cells than they are points in the cloud, NaNs and Infs are supported in a consistent manner: 
+- nans → random position
+- (+-) infs → border of the grid
+  
+This allow to gridsort prime or variable number of points, as long as one is ready to deal with void/special grid cells.
+
+Expected runtime for sorting 1 million points, tested on google colab GPU T4 for 2D/3D point clouds: 
+- Numpy: 10s
+- Torch (gpu): 500 ms
+- Cupy (gpu): 1 s (after compilation e.g. not on the first cold call, which take 10 s)
+
+`When not to use gridpoints ?`
+- high dimension: the package is implemented to support arbitrary dimension, but sweetspot is really 2D/3D. dimensions 4-6 might still be reasonable depending on the task, but anything above 8D is generally too high dimensional for gridpoints.
+- weird geometries. Supported datasets goes beyond smooth convex distributions: map of Indonesia, a sponge, a donuts, an elephant, an eggshell (by specifying a suitable gridshape e.g. (128,128,2) to `gridpoints.sort()`). You can see various exemples in the [plot](https://github.com/Neighborhood-Grid/CubeNet/blob/main/plots) folder But with some limits. Bad fits: a spider web, a wind turbine, same eggshell with a naive 3d sort (gridshape = (32,32,32)) ... the issue is not that gridpoints cannot sort these distributions, but that it will produce a poor representation of the geometry.
+- small point clouds: beyond a few hundred points, local operations in O(N) time is not worth the overhead, because naive quadratic implementations will probably be simultaneously simpler and faster.
+
+---
+
+### Installation
+
+```bash
+pip install gridpoints          # core only
+pip install gridpoints[demo]    # for the demonstration notebook.
+```
+Full Demo: [notebook.ipynb](https://github.com/Neighborhood-Grid/CubeNet/blob/main/notebook.ipynb)
+
+### Quickstart
+
+```python
+import gridpoints as grid
+import numpy as np
+
+# Raw point cloud (numpy, pytorch or cupy)
+A = np.random.rand(1_000_000, 3)
+
+# Sorted view: place the points inside the grid
+order = grid.argsort(A, gridshape=(100, 100, 100))
+Bflat = A[order] #flat but sorted with a grid layout (C-order)
+Bgrid = Bflat.reshape(100, 100, 100, 3) #reshaped as a grid 
+
+# Rest of your pipeline, working with grids
+Cgrid = apply_something(Bgrid) #(100, 100, 100, *)
+
+# Back to the original points indexing
+Cflat = Cgrid.reshape(1_000_000, *)
+orderinv = grid.invert_permutation(order)
+C = Cflat[orderinv]   # matches the initial points order
+```
+
+---
+
+### Note on rectangular gridshapes
+
+The eggshell example illustrates the importance of selecting a suitable grid-shape decomposition. This can itself be challenging when the point cloud is assumed to follow a complex topological pattern, such as a thin surface rather than a full volumetric domain. While the precise decomposition is not particularly important (e.g. (18,20,16) vs. (16,15,24) will generally make little difference), the orders of magnitude of the different dimensions do matter and should be tuned to the dataset: (18,20,16) and (36,40,4) can lead to substantially different results. As mentioned, nan/inf padding can help accomodate integer-factorization constraints
+
+Another possibility for smoothing out topology-specific effects is to perform the sort after randomly projecting the dataset onto a lower-dimensional subspace. Random projections are known to approximately preserve pairwise distances, as formalized by the Johnson–Lindenstrauss lemma.
+
+
+### Note on the cutoff radius R
+
+There is no strict theoretical guarantee about what the cutof radius R should be for a given task. E.g the relative grid position between a point and its nearest neighbors can't be garanted to be in the exact adjacent grid cells. What is guaranteed from the sorted ordering is only **grid monotonicity**: *x* coordinates increase along rows, *y* coordinates along columns, and so on.
+
+As an example, empirical results in 2-D show that `R = 5` is enough for ~99 % of the nearest neighbors; some outlier neighbors will sit further apart for complex geometries with pronounced peaks, holes or any non-smoothness. When a stricter neighborhood is required, or in high dimensional setting, the solution might be to build an assembly of grid experts, each working on a rotated / projected view of the points, as discussed in [this topic](https://github.com/glotzerlab/freud/discussions/1417). An other possibility would be to tile the grid: each tile is enhanced with location metadata (e.g. a bounding box), allowing for pruning pairs of tiles with a distance certified to be far enough for a given criterion.
+
+
+### Note on efficient stencil operations
+
+The typical use-case of `Gridpoints` is to allow fast local operations on arbitrary point clouds using stencil kernels:
+
+```text
+output(i, j, k) = f( Pgrid[i±di, j±dj, k±dk] | di, dj, dk in local window )
+```
+To go beyond standard (slow) python loops, this kind of kernel computation can be accelerated with native grid convolution operations of standard libraries whenever possible, or with `pystencils` or `taichi` compilers for complex/non O(N) grid kernels. On GPU, tiling the grid and leveraging a custom `triton` kernel might be particularly efficient.

{gridpoints-1.0.6 → gridpoints-1.1.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gridpoints"
-version = "1.0.6"
+version = "1.1.0"
 description = "fast multidimensional sorting tool for point clouds using the monotonic lagrangian grid "
 readme = "README.md"
 requires-python = ">=3.9"
@@ -26,8 +26,8 @@ demo = [
 
 
 [project.urls]
-Homepage = "https://github.com/Neighborhood-Grid/CubeNet"
-Repository = "https://github.com/Neighborhood-Grid/CubeNet"
+Homepage = "https://neighborhood-grid.github.io/gridpoints/"
+Repository = "https://github.com/Neighborhood-Grid/gridpoints"
 
 [tool.setuptools.packages.find]
 where = ["src"]

gridpoints-1.1.0/src/griddemo/__init__.py

@@ -0,0 +1,10 @@
+
+from .quick_demo import sample, show_result, print_zoom
+
+__all__ = (
+    "sample",
+    "show_result",
+    "print_zoom",
+)
+
+__version__ = "1.0.8"

{gridpoints-1.0.6 → gridpoints-1.1.0}/src/gridpoints/__init__.py

@@ -1,11 +1,13 @@
-"""Monotonic Lagrangian sorting for NumPy, CuPy, and PyTorch arrays."""
-
-from .mlg_sort import argsort, sort, invert_permutation
-
-__all__ = (
-    "argsort",
-    "sort",
-    "invert_permutation",
-)
-
-__version__ = "1.0.6"
+"""Monotonic Lagrangian sorting for NumPy, CuPy, and PyTorch arrays."""
+
+from .mlg_sort import argsort, sort, invert_permutation
+from .padding import pad
+
+__all__ = (
+    "argsort",
+    "sort",
+    "invert_permutation",
+    "pad",
+)
+
+__version__ = "1.1.0"

gridpoints-1.1.0/src/gridpoints/array_libs.py

@@ -0,0 +1,183 @@
+"""
+Helpers:
+Unified array syntax to support numpy, cupy and pytorch arrays
+
+Here we just define basic primitives that we will need later
+like arange, copy, shuffle etc.
+"""
+
+
+def precision(N):
+    """Names of the (float, int) dtypes suited to point clouds of N points.
+
+    - float32 resolves the gaps (~1/N) between neighbouring points only up to
+      a few 10**4 points, float64 is used beyond.
+    - the diagonal sorts store values up to 6N: int32 is safe up to N ~ 3.5e8.
+    """
+    return ('float32' if N <= 2**16 else 'float64'), ('int32' if 6 * N < 2**31 else 'int64')
+
+
+def get_dtype(dtp, xp):
+    if dtp == 'bool':
+        return getattr(xp, 'bool_', None) or xp.bool
+    return getattr(xp, dtp) if isinstance(dtp, str) else dtp
+
+
+def get_backend(X):
+    module = type(X).__module__.split('.')[0]
+    if module == 'torch':
+        import torch
+        return torch, 'torch'
+    elif module == 'cupy':
+        import cupy
+        return cupy, 'cupy'
+    else:
+        import numpy
+        return numpy, 'numpy'
+
+def astype(X, dtype_str, reference_array=None):
+    xp, backend_name = get_backend(X)
+    dtype = get_dtype(dtype_str, xp)
+    if backend_name == 'torch':
+        device = reference_array.device if reference_array is not None else X.device
+        return X.to(dtype=dtype, device=device)
+    return X.astype(dtype, copy=False)
+
+
+def inplace_sort(X, axis):
+    xp, backend_name = get_backend(X)
+    if backend_name == 'torch':
+        X.copy_(X.sort(dim=axis)[0])
+    elif backend_name == "cupy":
+        X[...] = xp.sort(xp.ascontiguousarray(X), axis=axis)
+    else:
+        X.sort(axis=axis)
+
+
+def sort(X, axis):
+    xp, backend_name = get_backend(X)
+    return X.sort(dim=axis)[0] if backend_name == 'torch' else xp.sort(X, axis=axis)
+
+
+def _alloc(fn, shape, dtype_str, reference_array):
+    xp, backend_name = get_backend(reference_array)
+    dtype = get_dtype(dtype_str, xp)
+    if backend_name == 'torch':
+        return getattr(xp, fn)(shape, dtype=dtype, device=reference_array.device)
+    return getattr(xp, fn)(shape, dtype=dtype)
+
+def empty(shape, dtype_str, reference_array):
+    return _alloc('empty', shape, dtype_str, reference_array)
+
+def zeros(shape, dtype_str, reference_array):
+    return _alloc('zeros', shape, dtype_str, reference_array)
+
+def ones(shape, dtype_str, reference_array):
+    return _alloc('ones', shape, dtype_str, reference_array)
+
+
+def arange(N, dtype_str, reference_array):
+    xp, backend_name = get_backend(reference_array)
+    dtype = get_dtype(dtype_str, xp)
+    if backend_name == 'torch':
+        return xp.arange(N, dtype=dtype, device=reference_array.device)
+    return xp.arange(N, dtype=dtype)
+
+def random(shape, dtype_str, reference_array):
+    """Uniform random numbers in [0, 1)."""
+    xp, backend_name = get_backend(reference_array)
+    dtype = get_dtype(dtype_str, xp)
+    if backend_name == 'torch':
+        return xp.rand(shape, dtype=dtype, device=reference_array.device)
+    elif backend_name == 'cupy':
+        return xp.random.rand(*shape).astype(dtype)
+    return xp.random.default_rng(42).random(shape, dtype=dtype)
+
+def copy(arr):
+    _, backend_name = get_backend(arr)
+    if backend_name == 'torch':
+        return arr.clone()
+    return arr.copy()
+
+def amax(X, axis):
+    return X.amax(axis) if get_backend(X)[1] == 'torch' else X.max(axis)
+
+def amin(X, axis):
+    return X.amin(axis) if get_backend(X)[1] == 'torch' else X.min(axis)
+
+def all_true(tensor):
+    return bool(tensor.all())
+
+def any_true(tensor):
+    return bool(tensor.any())
+
+
+def nan_infs_to_num_and_tiebreaker(arr):
+    """Make a batch of clouds (B, N, D) finite, and break ties with a tiny noise.
+
+    +inf / -inf are put just beyond the finite points of their cloud and axis
+    (they end up on the border of the grid), NaN are replaced by the value of a
+    random finite point of their cloud and axis (they end up anywhere).
+    """
+    xp, _ = get_backend(arr)
+    out = arr
+    finite = xp.isfinite(arr)
+    if not all_true(finite):
+        hi = amax(xp.where(finite, arr, -xp.inf), 1)[:, None, :]
+        lo = amin(xp.where(finite, arr, xp.inf), 1)[:, None, :]
+        out = xp.where(xp.isposinf(arr), hi + 1, arr)
+        out = xp.where(xp.isneginf(arr), lo - 1, out)
+        nan = xp.isnan(arr)
+        if any_true(nan):
+            finite_first = sort(xp.where(finite, arr, xp.inf), 1)
+            n_finite = finite.sum(1)[:, None, :]
+            pick = astype(random(arr.shape, arr.dtype, arr) * n_finite, 'int64')
+            out = xp.where(nan, take_along_axis(finite_first, pick, 1), out)
+    eps = 1e-5 if '32' in str(arr.dtype) else 1e-10
+    return out + eps * random(out.shape, arr.dtype, out)
+
+def flip(X, axis=None):
+    xp, _ = get_backend(X)
+    if axis is None or axis == () or axis == []:
+        return X
+
+    if isinstance(axis, int):
+        axis = (axis,)
+    else:
+        axis = tuple(axis)
+
+    if hasattr(X, "flip") and callable(X.flip):
+        return X.flip(axis)
+
+    return xp.flip(X, axis=axis)
+
+def argsort(X):
+    xp, backend_name = get_backend(X)
+    return xp.argsort(X, dim=-1) if backend_name == 'torch' else xp.argsort(X, axis=-1)
+
+def argpartition(X, mid):
+    """Indices that put the `mid` smallest values of the last axis first."""
+    xp, backend_name = get_backend(X)
+    if backend_name == 'torch':
+        return xp.argsort(X, dim=-1)
+    else:
+        return xp.argpartition(X, mid - 1, axis=-1)
+
+def take_along_axis(arr, idx, axis):
+    xp, backend_name = get_backend(arr)
+    if backend_name == 'torch':
+        return xp.take_along_dim(arr, idx.long(), dim=axis)
+    else:
+        return xp.take_along_axis(arr, idx, axis=axis)
+
+def concatenate(arrays, axis = 0):
+    xp, name = get_backend(arrays[0])
+    if name == "torch":
+        return xp.cat(list(arrays), dim=axis)
+    return xp.concatenate(arrays, axis=axis)
+
+def transpose(X, axes):
+    _, name = get_backend(X)
+    if name == "torch":
+        return X.permute(*axes)
+    return X.transpose(axes)

gridpoints-1.1.0/src/gridpoints/diagonal_tricks.py

@@ -0,0 +1,106 @@
+"""
+Helpers:
+Diagonal trick is based on a tensor rotation that allow
+to sort arrays efficiently along arbitrary diagonal fibers.
+
+All functions work on a batch of grids X of shape (B, ...), and tell for each
+grid of the batch whether it was already sorted.
+"""
+
+import math
+
+from .array_libs import (
+    get_backend,
+    arange,
+    inplace_sort,
+)
+
+
+def sorted_flags(X, axis):
+    """Bool array of shape (B,): is each grid of the batch X sorted along `axis`?"""
+    axis %= X.ndim
+    sl = lambda s: (slice(None),) * axis + (s,)
+    ok = X[sl(slice(1, None))] >= X[sl(slice(None, -1))]
+    return ok.reshape(len(X), math.prod(ok.shape[1:])).all(1)
+
+
+def sort_diag_trick(X):
+    """
+    we rotate the tensor along the main diagonal
+    to place diagonals on the last axis, then
+    perform a standard sort(axis = -1), and
+    rotate back the diagonals
+    """
+    xp, _ = get_backend(X)
+    #bigM for safety, protect
+    #overlapping diagonals
+    bigM = (X.max() - X.min()) + 1
+
+    n, m = X.shape[-2:]
+
+    #move the smallest axis to last position
+    if m > n:
+        X = xp.moveaxis(X, [-1, -2], [-2, -1])
+        n, m = m, n #now n >= m
+
+    i = arange(n, X.dtype, X)[:, None]
+    j = arange(m, X.dtype, X)[None, :]
+
+    di = (i - j)
+    i_ = di % n
+    ires = di // n
+    groupkey = bigM * ires
+
+    X[..., i_, j] = X[..., i, j] - groupkey #Rotation
+    # ========================
+    # here is the trick: as we rotated the tensor
+    # sorting on the diagonal is just working on
+    # axis -1 of the rotated view
+    # ========================
+    was_sorted = sorted_flags(X, -1)
+    inplace_sort(X, axis=-1)
+    X[..., i, j] = X[..., i_, j] + groupkey
+
+    return was_sorted
+
+def sort_tridiag_trick(X):
+    """
+    we rotate the tensor along the tridiagonal,
+    then perform a standard sort(axis = -1),
+    and rotate back the tridiagonals
+    """
+    xp, _ = get_backend(X)
+    #bigM for safety, protect
+    #overlapping diagonals
+    bigM = (X.max() - X.min()) + 1
+
+    #move the smallest axes to last positions
+    abc = X.shape
+    dst_axes = (-3, -2, -1)
+    src_axes = tuple(sorted(dst_axes, key=lambda ax: -abc[ax]))
+    X = xp.moveaxis(X, src_axes, dst_axes)
+    a, b, c = X.shape[-3:]  # now a >= b >= c
+
+    i = arange(a, X.dtype, X)[:, None, None]
+    j = arange(b, X.dtype, X)[None, :, None]
+    k = arange(c, X.dtype, X)[None, None, :]
+
+    di = i - k
+    dj = j - k
+    di_ = di % a
+    dj_ = dj % b
+    dires = di // a
+    djres = dj // b
+    groupkey = bigM * (dires * 4 + djres)
+
+    X[..., di_, dj_, k] = X[..., i, j, k] - groupkey #Trirotation
+    # ========================
+    # Again, thanks to the rotation trick
+    # we just need to .sort(X, axis = -1)
+    # to sort the tridiagonal
+    # ========================
+    was_sorted = sorted_flags(X, -1)
+    inplace_sort(X, axis=-1)
+    X[..., i, j, k] = X[..., di_, dj_, k] + groupkey
+
+    return was_sorted