ocnn 2.2.1__py3-none-any.whl → 2.2.3__py3-none-any.whl

ocnn/modules/resblocks.py

@@ -1,124 +1,124 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-import torch
-import torch.utils.checkpoint
-
-from ocnn.octree import Octree
-from ocnn.nn import OctreeMaxPool
-from ocnn.modules import Conv1x1BnRelu, OctreeConvBnRelu, Conv1x1Bn, OctreeConvBn
-
-
-class OctreeResBlock(torch.nn.Module):
-  r''' Octree-based ResNet block in a bottleneck style. The block is composed of
-  a series of :obj:`Conv1x1`, :obj:`Conv3x3`, and :obj:`Conv1x1`.
-
-  Args:
-    in_channels (int): Number of input channels.
-    out_channels (int): Number of output channels.
-    stride (int): The stride of the block (:obj:`1` or :obj:`2`).
-    bottleneck (int): The input and output channels of the :obj:`Conv3x3` is
-        equal to the input channel divided by :attr:`bottleneck`.
-    nempty (bool): If True, only performs the convolution on non-empty
-        octree nodes.
-  '''
-
-  def __init__(self, in_channels: int, out_channels: int, stride: int = 1,
-               bottleneck: int = 4, nempty: bool = False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.bottleneck = bottleneck
-    self.stride = stride
-    channelb = int(out_channels / bottleneck)
-
-    if self.stride == 2:
-      self.max_pool = OctreeMaxPool(nempty)
-    self.conv1x1a = Conv1x1BnRelu(in_channels, channelb)
-    self.conv3x3 = OctreeConvBnRelu(channelb, channelb, nempty=nempty)
-    self.conv1x1b = Conv1x1Bn(channelb, out_channels)
-    if self.in_channels != self.out_channels:
-      self.conv1x1c = Conv1x1Bn(in_channels, out_channels)
-    self.relu = torch.nn.ReLU(inplace=True)
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    if self.stride == 2:
-      data = self.max_pool(data, octree, depth)
-      depth = depth - 1
-    conv1 = self.conv1x1a(data)
-    conv2 = self.conv3x3(conv1, octree, depth)
-    conv3 = self.conv1x1b(conv2)
-    if self.in_channels != self.out_channels:
-      data = self.conv1x1c(data)
-    out = self.relu(conv3 + data)
-    return out
-
-
-class OctreeResBlock2(torch.nn.Module):
-  r''' Basic Octree-based ResNet block. The block is composed of
-  a series of :obj:`Conv3x3` and :obj:`Conv3x3`.
-
-  Refer to :class:`OctreeResBlock` for the details of arguments.
-  '''
-
-  def __init__(self, in_channels, out_channels, stride=1, bottleneck=1,
-               nempty=False):
-    super().__init__()
-    self.in_channels = in_channels
-    self.out_channels = out_channels
-    self.stride = stride
-    channelb = int(out_channels / bottleneck)
-
-    if self.stride == 2:
-      self.maxpool = OctreeMaxPool(self.depth)
-    self.conv3x3a = OctreeConvBnRelu(in_channels, channelb, nempty=nempty)
-    self.conv3x3b = OctreeConvBn(channelb, out_channels, nempty=nempty)
-    if self.in_channels != self.out_channels:
-      self.conv1x1 = Conv1x1Bn(in_channels, out_channels)
-    self.relu = torch.nn.ReLU(inplace=True)
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    if self.stride == 2:
-      data = self.maxpool(data, octree, depth)
-      depth = depth - 1
-    conv1 = self.conv3x3a(data, octree, depth)
-    conv2 = self.conv3x3b(conv1, octree, depth)
-    if self.in_channels != self.out_channels:
-      data = self.conv1x1(data)
-    out = self.relu(conv2 + data)
-    return out
-
-
-class OctreeResBlocks(torch.nn.Module):
-  r''' A sequence of :attr:`resblk_num` ResNet blocks.
-  '''
-
-  def __init__(self, in_channels, out_channels, resblk_num, bottleneck=4,
-               nempty=False, resblk=OctreeResBlock, use_checkpoint=False):
-    super().__init__()
-    self.resblk_num = resblk_num
-    self.use_checkpoint = use_checkpoint
-    channels = [in_channels] + [out_channels] * resblk_num
-
-    self.resblks = torch.nn.ModuleList(
-        [resblk(channels[i], channels[i+1], 1, bottleneck, nempty)
-         for i in range(self.resblk_num)])
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    for i in range(self.resblk_num):
-      if self.use_checkpoint:
-        data = torch.utils.checkpoint.checkpoint(
-            self.resblks[i], data, octree, depth)
-      else:
-        data = self.resblks[i](data, octree, depth)
-    return data
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+import torch.utils.checkpoint
+
+from ocnn.octree import Octree
+from ocnn.nn import OctreeMaxPool
+from ocnn.modules import Conv1x1BnRelu, OctreeConvBnRelu, Conv1x1Bn, OctreeConvBn
+
+
+class OctreeResBlock(torch.nn.Module):
+  r''' Octree-based ResNet block in a bottleneck style. The block is composed of
+  a series of :obj:`Conv1x1`, :obj:`Conv3x3`, and :obj:`Conv1x1`.
+
+  Args:
+    in_channels (int): Number of input channels.
+    out_channels (int): Number of output channels.
+    stride (int): The stride of the block (:obj:`1` or :obj:`2`).
+    bottleneck (int): The input and output channels of the :obj:`Conv3x3` is
+        equal to the input channel divided by :attr:`bottleneck`.
+    nempty (bool): If True, only performs the convolution on non-empty
+        octree nodes.
+  '''
+
+  def __init__(self, in_channels: int, out_channels: int, stride: int = 1,
+               bottleneck: int = 4, nempty: bool = False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.bottleneck = bottleneck
+    self.stride = stride
+    channelb = int(out_channels / bottleneck)
+
+    if self.stride == 2:
+      self.max_pool = OctreeMaxPool(nempty)
+    self.conv1x1a = Conv1x1BnRelu(in_channels, channelb)
+    self.conv3x3 = OctreeConvBnRelu(channelb, channelb, nempty=nempty)
+    self.conv1x1b = Conv1x1Bn(channelb, out_channels)
+    if self.in_channels != self.out_channels:
+      self.conv1x1c = Conv1x1Bn(in_channels, out_channels)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    if self.stride == 2:
+      data = self.max_pool(data, octree, depth)
+      depth = depth - 1
+    conv1 = self.conv1x1a(data)
+    conv2 = self.conv3x3(conv1, octree, depth)
+    conv3 = self.conv1x1b(conv2)
+    if self.in_channels != self.out_channels:
+      data = self.conv1x1c(data)
+    out = self.relu(conv3 + data)
+    return out
+
+
+class OctreeResBlock2(torch.nn.Module):
+  r''' Basic Octree-based ResNet block. The block is composed of
+  a series of :obj:`Conv3x3` and :obj:`Conv3x3`.
+
+  Refer to :class:`OctreeResBlock` for the details of arguments.
+  '''
+
+  def __init__(self, in_channels, out_channels, stride=1, bottleneck=1,
+               nempty=False):
+    super().__init__()
+    self.in_channels = in_channels
+    self.out_channels = out_channels
+    self.stride = stride
+    channelb = int(out_channels / bottleneck)
+
+    if self.stride == 2:
+      self.maxpool = OctreeMaxPool(self.depth)
+    self.conv3x3a = OctreeConvBnRelu(in_channels, channelb, nempty=nempty)
+    self.conv3x3b = OctreeConvBn(channelb, out_channels, nempty=nempty)
+    if self.in_channels != self.out_channels:
+      self.conv1x1 = Conv1x1Bn(in_channels, out_channels)
+    self.relu = torch.nn.ReLU(inplace=True)
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    if self.stride == 2:
+      data = self.maxpool(data, octree, depth)
+      depth = depth - 1
+    conv1 = self.conv3x3a(data, octree, depth)
+    conv2 = self.conv3x3b(conv1, octree, depth)
+    if self.in_channels != self.out_channels:
+      data = self.conv1x1(data)
+    out = self.relu(conv2 + data)
+    return out
+
+
+class OctreeResBlocks(torch.nn.Module):
+  r''' A sequence of :attr:`resblk_num` ResNet blocks.
+  '''
+
+  def __init__(self, in_channels, out_channels, resblk_num, bottleneck=4,
+               nempty=False, resblk=OctreeResBlock, use_checkpoint=False):
+    super().__init__()
+    self.resblk_num = resblk_num
+    self.use_checkpoint = use_checkpoint
+    channels = [in_channels] + [out_channels] * resblk_num
+
+    self.resblks = torch.nn.ModuleList(
+        [resblk(channels[i], channels[i+1], 1, bottleneck, nempty)
+         for i in range(self.resblk_num)])
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    for i in range(self.resblk_num):
+      if self.use_checkpoint:
+        data = torch.utils.checkpoint.checkpoint(
+            self.resblks[i], data, octree, depth)
+      else:
+        data = self.resblks[i](data, octree, depth)
+    return data

ocnn/nn/__init__.py

@@ -1,42 +1,43 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-from .octree2vox import octree2voxel, Octree2Voxel
-from .octree2col import octree2col, col2octree
-from .octree_pad import octree_pad, octree_depad
-from .octree_interp import (octree_nearest_pts, octree_linear_pts,
-                            OctreeInterp, OctreeUpsample)
-from .octree_pool import (octree_max_pool, OctreeMaxPool,
-                          octree_max_unpool, OctreeMaxUnpool,
-                          octree_global_pool, OctreeGlobalPool,
-                          octree_avg_pool, OctreeAvgPool,)
-from .octree_conv import OctreeConv, OctreeDeconv
-from .octree_dwconv import OctreeDWConv
-from .octree_norm import OctreeBatchNorm, OctreeGroupNorm, OctreeInstanceNorm
-from .octree_drop import OctreeDropPath
-from .octree_align import search_value, octree_align
-
-
-__all__ = [
-    'octree2voxel',
-    'octree2col', 'col2octree',
-    'octree_pad', 'octree_depad',
-    'octree_nearest_pts', 'octree_linear_pts',
-    'octree_max_pool', 'octree_max_unpool',
-    'octree_global_pool', 'octree_avg_pool',
-    'Octree2Voxel',
-    'OctreeMaxPool', 'OctreeMaxUnpool',
-    'OctreeGlobalPool', 'OctreeAvgPool',
-    'OctreeConv', 'OctreeDeconv',
-    'OctreeDWConv',
-    'OctreeInterp', 'OctreeUpsample',
-    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm',
-    'OctreeDropPath',
-    'search_value', 'octree_align',
-]
-
-classes = __all__
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+from .octree2vox import octree2voxel, Octree2Voxel
+from .octree2col import octree2col, col2octree
+from .octree_pad import octree_pad, octree_depad
+from .octree_interp import (octree_nearest_pts, octree_linear_pts,
+                            OctreeInterp, OctreeUpsample)
+from .octree_pool import (octree_max_pool, OctreeMaxPool,
+                          octree_max_unpool, OctreeMaxUnpool,
+                          octree_global_pool, OctreeGlobalPool,
+                          octree_avg_pool, OctreeAvgPool,)
+from .octree_conv import OctreeConv, OctreeDeconv
+from .octree_gconv import OctreeGroupConv
+from .octree_dwconv import OctreeDWConv
+from .octree_norm import OctreeBatchNorm, OctreeGroupNorm, OctreeInstanceNorm
+from .octree_drop import OctreeDropPath
+from .octree_align import search_value, octree_align
+
+
+__all__ = [
+    'octree2voxel',
+    'octree2col', 'col2octree',
+    'octree_pad', 'octree_depad',
+    'octree_nearest_pts', 'octree_linear_pts',
+    'octree_max_pool', 'octree_max_unpool',
+    'octree_global_pool', 'octree_avg_pool',
+    'Octree2Voxel',
+    'OctreeMaxPool', 'OctreeMaxUnpool',
+    'OctreeGlobalPool', 'OctreeAvgPool',
+    'OctreeConv', 'OctreeDeconv',
+    'OctreeGroupConv', 'OctreeDWConv',
+    'OctreeInterp', 'OctreeUpsample',
+    'OctreeInstanceNorm', 'OctreeBatchNorm', 'OctreeGroupNorm',
+    'OctreeDropPath',
+    'search_value', 'octree_align',
+]
+
+classes = __all__

ocnn/nn/octree2col.py

@@ -1,53 +1,53 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-import torch
-import torch.nn
-
-from ocnn.octree import Octree
-from ocnn.utils import scatter_add
-
-
-def octree2col(data: torch.Tensor, octree: Octree, depth: int,
-               kernel_size: str = '333', stride: int = 1, nempty: bool = False):
-  r''' Gathers the neighboring features for convolutions.
-
-  Args:
-    data (torch.Tensor): The input data.
-    octree (Octree): The corresponding octree.
-    depth (int): The depth of current octree.
-    kernel_size (str): The kernel shape, choose from :obj:`333`, :obj:`311`, 
-        :obj:`131`, :obj:`113`, :obj:`222`, :obj:`331`, :obj:`133`, and
-        :obj:`313`.
-    stride (int): The stride of neighborhoods (:obj:`1` or :obj:`2`). If the
-        stride is :obj:`2`, it always returns the neighborhood of the first
-        siblings, and the number of elements of output tensor is
-        :obj:`octree.nnum[depth] / 8`.
-    nempty (bool): If True, only returns the neighborhoods of the non-empty
-        octree nodes.
-  '''
-
-  neigh = octree.get_neigh(depth, kernel_size, stride, nempty)
-  size = (neigh.shape[0], neigh.shape[1], data.shape[1])
-  out = torch.zeros(size, dtype=data.dtype, device=data.device)
-  valid = neigh >= 0
-  out[valid] = data[neigh[valid]]  # (N, K, C)
-  return out
-
-
-def col2octree(data: torch.Tensor, octree: Octree, depth: int,
-               kernel_size: str = '333', stride: int = 1, nempty: bool = False):
-  r''' Scatters the convolution features to an octree.
-
-  Please refer to :func:`octree2col` for the usage of function parameters.
-  '''
-
-  neigh = octree.get_neigh(depth, kernel_size, stride, nempty)
-  valid = neigh >= 0
-  dim_size = octree.nnum_nempty[depth] if nempty else octree.nnum[depth]
-  out = scatter_add(data[valid], neigh[valid], dim=0, dim_size=dim_size)
-  return out
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+import torch.nn
+
+from ocnn.octree import Octree
+from ocnn.utils import scatter_add
+
+
+def octree2col(data: torch.Tensor, octree: Octree, depth: int,
+               kernel_size: str = '333', stride: int = 1, nempty: bool = False):
+  r''' Gathers the neighboring features for convolutions.
+
+  Args:
+    data (torch.Tensor): The input data.
+    octree (Octree): The corresponding octree.
+    depth (int): The depth of current octree.
+    kernel_size (str): The kernel shape, choose from :obj:`333`, :obj:`311`, 
+        :obj:`131`, :obj:`113`, :obj:`222`, :obj:`331`, :obj:`133`, and
+        :obj:`313`.
+    stride (int): The stride of neighborhoods (:obj:`1` or :obj:`2`). If the
+        stride is :obj:`2`, it always returns the neighborhood of the first
+        siblings, and the number of elements of output tensor is
+        :obj:`octree.nnum[depth] / 8`.
+    nempty (bool): If True, only returns the neighborhoods of the non-empty
+        octree nodes.
+  '''
+
+  neigh = octree.get_neigh(depth, kernel_size, stride, nempty)
+  size = (neigh.shape[0], neigh.shape[1], data.shape[1])
+  out = torch.zeros(size, dtype=data.dtype, device=data.device)
+  valid = neigh >= 0
+  out[valid] = data[neigh[valid]]  # (N, K, C)
+  return out
+
+
+def col2octree(data: torch.Tensor, octree: Octree, depth: int,
+               kernel_size: str = '333', stride: int = 1, nempty: bool = False):
+  r''' Scatters the convolution features to an octree.
+
+  Please refer to :func:`octree2col` for the usage of function parameters.
+  '''
+
+  neigh = octree.get_neigh(depth, kernel_size, stride, nempty)
+  valid = neigh >= 0
+  dim_size = octree.nnum_nempty[depth] if nempty else octree.nnum[depth]
+  out = scatter_add(data[valid], neigh[valid], dim=0, dim_size=dim_size)
+  return out

ocnn/nn/octree2vox.py

@@ -1,50 +1,50 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-import torch
-
-from ocnn.octree import Octree
-
-
-def octree2voxel(data: torch.Tensor, octree: Octree, depth: int,
-                 nempty: bool = False):
-  r''' Converts the input feature to the full-voxel-based representation.
-
-  Args:
-    data (torch.Tensor): The input feature.
-    octree (Octree): The corresponding octree.
-    depth (int): The depth of current octree.
-    nempty (bool): If True, :attr:`data` only contains the features of non-empty
-        octree nodes.
-  '''
-
-  x, y, z, b = octree.xyzb(depth, nempty)
-
-  num = 1 << depth
-  channel = data.shape[1]
-  vox = data.new_zeros([octree.batch_size, num, num, num, channel])
-  vox[b, x, y, z] = data
-  return vox
-
-
-class Octree2Voxel(torch.nn.Module):
-  r''' Converts the input feature to the full-voxel-based representation.
-
-  Please refer to :func:`octree2voxel` for details.
-  '''
-
-  def __init__(self, nempty: bool = False):
-    super().__init__()
-    self.nempty = nempty
-
-  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
-    r''''''
-
-    return octree2voxel(data, octree, depth, self.nempty)
-
-  def extra_repr(self) -> str:
-    return 'nempty={}'.format(self.nempty)
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+
+from ocnn.octree import Octree
+
+
+def octree2voxel(data: torch.Tensor, octree: Octree, depth: int,
+                 nempty: bool = False):
+  r''' Converts the input feature to the full-voxel-based representation.
+
+  Args:
+    data (torch.Tensor): The input feature.
+    octree (Octree): The corresponding octree.
+    depth (int): The depth of current octree.
+    nempty (bool): If True, :attr:`data` only contains the features of non-empty
+        octree nodes.
+  '''
+
+  x, y, z, b = octree.xyzb(depth, nempty)
+
+  num = 1 << depth
+  channel = data.shape[1]
+  vox = data.new_zeros([octree.batch_size, num, num, num, channel])
+  vox[b, x, y, z] = data
+  return vox
+
+
+class Octree2Voxel(torch.nn.Module):
+  r''' Converts the input feature to the full-voxel-based representation.
+
+  Please refer to :func:`octree2voxel` for details.
+  '''
+
+  def __init__(self, nempty: bool = False):
+    super().__init__()
+    self.nempty = nempty
+
+  def forward(self, data: torch.Tensor, octree: Octree, depth: int):
+    r''''''
+
+    return octree2voxel(data, octree, depth, self.nempty)
+
+  def extra_repr(self) -> str:
+    return 'nempty={}'.format(self.nempty)

ocnn/nn/octree_align.py

@@ -1,46 +1,46 @@
-# --------------------------------------------------------
-# Octree-based Sparse Convolutional Neural Networks
-# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
-# Licensed under The MIT License [see LICENSE for details]
-# Written by Peng-Shuai Wang
-# --------------------------------------------------------
-
-import torch
-
-from ocnn.octree import Octree
-
-
-def search_value(value: torch.Tensor, key: torch.Tensor, query: torch.Tensor):
-  r''' Searches values according to sorted shuffled keys.
-
-  Args:
-    value (torch.Tensor): The input tensor with shape (N, C).
-    key (torch.Tensor): The key tensor corresponds to :attr:`value` with shape 
-        (N,), which contains sorted shuffled keys of an octree.
-    query (torch.Tensor): The query tensor, which also contains shuffled keys.
-  '''
-
-  # deal with out-of-bound queries, the indices of these queries
-  # returned by torch.searchsorted equal to `key.shape[0]`
-  out_of_bound = query > key[-1]
-
-  # search
-  idx = torch.searchsorted(key, query)
-  idx[out_of_bound] = -1   # to avoid overflow when executing the following line
-  found = key[idx] == query
-
-  # assign the found value to the output
-  out = torch.zeros(query.shape[0], value.shape[1], device=value.device)
-  out[found] = value[idx[found]]
-  return out
-
-
-def octree_align(value: torch.Tensor, octree: Octree, octree_query: Octree,
-                 depth: int, nempty: bool = False):
-  r''' Wraps :func:`octree_align` to take octrees as input for convenience.
-  '''
-
-  key = octree.key(depth, nempty)
-  query = octree_query.key(depth, nempty)
-  assert key.shape[0] == value.shape[0]
-  return search_value(value, key, query)
+# --------------------------------------------------------
+# Octree-based Sparse Convolutional Neural Networks
+# Copyright (c) 2022 Peng-Shuai Wang <wangps@hotmail.com>
+# Licensed under The MIT License [see LICENSE for details]
+# Written by Peng-Shuai Wang
+# --------------------------------------------------------
+
+import torch
+
+from ocnn.octree import Octree
+
+
+def search_value(value: torch.Tensor, key: torch.Tensor, query: torch.Tensor):
+  r''' Searches values according to sorted shuffled keys.
+
+  Args:
+    value (torch.Tensor): The input tensor with shape (N, C).
+    key (torch.Tensor): The key tensor corresponds to :attr:`value` with shape 
+        (N,), which contains sorted shuffled keys of an octree.
+    query (torch.Tensor): The query tensor, which also contains shuffled keys.
+  '''
+
+  # deal with out-of-bound queries, the indices of these queries
+  # returned by torch.searchsorted equal to `key.shape[0]`
+  out_of_bound = query > key[-1]
+
+  # search
+  idx = torch.searchsorted(key, query)
+  idx[out_of_bound] = -1   # to avoid overflow when executing the following line
+  found = key[idx] == query
+
+  # assign the found value to the output
+  out = torch.zeros(query.shape[0], value.shape[1], device=value.device)
+  out[found] = value[idx[found]]
+  return out
+
+
+def octree_align(value: torch.Tensor, octree: Octree, octree_query: Octree,
+                 depth: int, nempty: bool = False):
+  r''' Wraps :func:`octree_align` to take octrees as input for convenience.
+  '''
+
+  key = octree.key(depth, nempty)
+  query = octree_query.key(depth, nempty)
+  assert key.shape[0] == value.shape[0]
+  return search_value(value, key, query)