content: from typing import *
  import torch
  import torch.nn as nn
  
  from trellis.backend_config import get_debug_mode, get_spconv_algo, get_sparse_backend
  
  
  SparseTensorData = None # Lazy import
  
  DEBUG = get_debug_mode()
  
  SPCONV_ALGO = get_spconv_algo()
  
  BACKEND = get_sparse_backend()
  __all__ = [
      'SparseTensor',
      'sparse_batch_broadcast',
      'sparse_batch_op',
      'sparse_cat',
      'sparse_unbind',
  ]
  
  
  class SparseTensor:
      """
      Sparse tensor with support for both torchsparse and spconv backends.
      
      Parameters:
      - feats (torch.Tensor): Features of the sparse tensor.
      - coords (torch.Tensor): Coordinates of the sparse tensor.
      - shape (torch.Size): Shape of the sparse tensor.
      - layout (List[slice]): Layout of the sparse tensor for each batch
      - data (SparseTensorData): Sparse tensor data used for convolusion
  
      NOTE:
      - Data corresponding to a same batch should be contiguous.
      - Coords should be in [0, 1023]
      """
      @overload
      def __init__(self, feats: torch.Tensor, coords: torch.Tensor, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ...
  
      @overload
      def __init__(self, data, shape: Optional[torch.Size] = None, layout: Optional[List[slice]] = None, **kwargs): ...
  
      def __init__(self, *args, **kwargs):
          # Lazy import of sparse tensor backend
          global SparseTensorData
          if SparseTensorData is None:
              import importlib
              if BACKEND == 'torchsparse':
                  SparseTensorData = importlib.import_module('torchsparse').SparseTensor
              elif BACKEND == 'spconv':
                  SparseTensorData = importlib.import_module('spconv.pytorch').SparseConvTensor
                  
          method_id = 0
          if len(args) != 0:
              method_id = 0 if isinstance(args[0], torch.Tensor) else 1
          else:
              method_id = 1 if 'data' in kwargs else 0
  
          if method_id == 0:
              feats, coords, shape, layout = args + (None,) * (4 - len(args))
              if 'feats' in kwargs:
                  feats = kwargs['feats']
                  del kwargs['feats']
              if 'coords' in kwargs:
                  coords = kwargs['coords']
                  del kwargs['coords']
              if 'shape' in kwargs:
                  shape = kwargs['shape']
                  del kwargs['shape']
              if 'layout' in kwargs:
                  layout = kwargs['layout']
                  del kwargs['layout']
  
              if shape is None:
                  shape = self.__cal_shape(feats, coords)
              if layout is None:
                  layout = self.__cal_layout(coords, shape[0])
              if BACKEND == 'torchsparse':
                  self.data = SparseTensorData(feats, coords, **kwargs)
              elif BACKEND == 'spconv':
                  spatial_shape = list(coords.max(0)[0] + 1)[1:]
                  self.data = SparseTensorData(feats.reshape(feats.shape[0], -1), coords, spatial_shape, shape[0], **kwargs)
                  self.data._features = feats
          elif method_id == 1:
              data, shape, layout = args + (None,) * (3 - len(args))
              if 'data' in kwargs:
                  data = kwargs['data']
                  del kwargs['data']
              if 'shape' in kwargs:
                  shape = kwargs['shape']
                  del kwargs['shape']
              if 'layout' in kwargs:
                  layout = kwargs['layout']
                  del kwargs['layout']
  
              self.data = data
              if shape is None:
                  shape = self.__cal_shape(self.feats, self.coords)
              if layout is None:
                  layout = self.__cal_layout(self.coords, shape[0])
  
          self._shape = shape
          self._layout = layout
          self._scale = kwargs.get('scale', (1, 1, 1))
          self._spatial_cache = kwargs.get('spatial_cache', {})
  
          if DEBUG:
              try:
                  assert self.feats.shape[0] == self.coords.shape[0], f"Invalid feats shape: {self.feats.shape}, coords shape: {self.coords.shape}"
                  assert self.shape == self.__cal_shape(self.feats, self.coords), f"Invalid shape: {self.shape}"
                  assert self.layout == self.__cal_layout(self.coords, self.shape[0]), f"Invalid layout: {self.layout}"
                  for i in range(self.shape[0]):
                      assert torch.all(self.coords[self.layout[i], 0] == i), f"The data of batch {i} is not contiguous"
              except Exception as e:
                  print('Debugging information:')
                  print(f"- Shape: {self.shape}")
                  print(f"- Layout: {self.layout}")
                  print(f"- Scale: {self._scale}")
                  print(f"- Coords: {self.coords}")
                  raise e
          
      def __cal_shape(self, feats, coords):
          shape = []
          shape.append(coords[:, 0].max().item() + 1)
          shape.extend([*feats.shape[1:]])
          return torch.Size(shape)
      
      def __cal_layout(self, coords, batch_size):
          seq_len = torch.bincount(coords[:, 0], minlength=batch_size)
          offset = torch.cumsum(seq_len, dim=0) 
          layout = [slice((offset[i] - seq_len[i]).item(), offset[i].item()) for i in range(batch_size)]
          return layout
      
      @property
      def shape(self) -> torch.Size:
          return self._shape
      
      def dim(self) -> int:
          return len(self.shape)
      
      @property
      def layout(self) -> List[slice]:
          return self._layout
  
      @property
      def feats(self) -> torch.Tensor:
          if BACKEND == 'torchsparse':
              return self.data.F
          elif BACKEND == 'spconv':
              return self.data.features
      
      @feats.setter
      def feats(self, value: torch.Tensor):
          if BACKEND == 'torchsparse':
              self.data.F = value
          elif BACKEND == 'spconv':
              self.data.features = value
  
      @property
      def coords(self) -> torch.Tensor:
          if BACKEND == 'torchsparse':
              return self.data.C
          elif BACKEND == 'spconv':
              return self.data.indices
          
      @coords.setter
      def coords(self, value: torch.Tensor):
          if BACKEND == 'torchsparse':
              self.data.C = value
          elif BACKEND == 'spconv':
              self.data.indices = value
  
      @property
      def dtype(self):
          return self.feats.dtype
  
      @property
      def device(self):
          return self.feats.device
  
      @overload
      def to(self, dtype: torch.dtype) -> 'SparseTensor': ...
  
      @overload
      def to(self, device: Optional[Union[str, torch.device]] = None, dtype: Optional[torch.dtype] = None) -> 'SparseTensor': ...
  
      def to(self, *args, **kwargs) -> 'SparseTensor':
          device = None
          dtype = None
          if len(args) == 2:
              device, dtype = args
          elif len(args) == 1:
              if isinstance(args[0], torch.dtype):
                  dtype = args[0]
              else:
                  device = args[0]
          if 'dtype' in kwargs:
              assert dtype is None, "to() received multiple values for argument 'dtype'"
              dtype = kwargs['dtype']
          if 'device' in kwargs:
              assert device is None, "to() received multiple values for argument 'device'"
              device = kwargs['device']
          
          new_feats = self.feats.to(device=device, dtype=dtype)
          new_coords = self.coords.to(device=device)
          return self.replace(new_feats, new_coords)
  
      def type(self, dtype):
          new_feats = self.feats.type(dtype)
          return self.replace(new_feats)
  
      def cpu(self) -> 'SparseTensor':
          new_feats = self.feats.cpu()
          new_coords = self.coords.cpu()
          return self.replace(new_feats, new_coords)
      
      def cuda(self) -> 'SparseTensor':
          new_feats = self.feats.cuda()
          new_coords = self.coords.cuda()
          return self.replace(new_feats, new_coords)
  
      def half(self) -> 'SparseTensor':
          new_feats = self.feats.half()
          return self.replace(new_feats)
      
      def float(self) -> 'SparseTensor':
          new_feats = self.feats.float()
          return self.replace(new_feats)
      
      def detach(self) -> 'SparseTensor':
          new_coords = self.coords.detach()
          new_feats = self.feats.detach()
          return self.replace(new_feats, new_coords)
  
      def dense(self) -> torch.Tensor:
          if BACKEND == 'torchsparse':
              return self.data.dense()
          elif BACKEND == 'spconv':
              return self.data.dense()
  
      def reshape(self, *shape) -> 'SparseTensor':
          new_feats = self.feats.reshape(self.feats.shape[0], *shape)
          return self.replace(new_feats)
      
      def unbind(self, dim: int) -> List['SparseTensor']:
          return sparse_unbind(self, dim)
  
      def replace(self, feats: torch.Tensor, coords: Optional[torch.Tensor] = None) -> 'SparseTensor':
          new_shape = [self.shape[0]]
          new_shape.extend(feats.shape[1:])
          if BACKEND == 'torchsparse':
              new_data = SparseTensorData(
                  feats=feats,
                  coords=self.data.coords if coords is None else coords,
                  stride=self.data.stride,
                  spatial_range=self.data.spatial_range,
              )
              new_data._caches = self.data._caches
          elif BACKEND == 'spconv':
              new_data = SparseTensorData(
                  self.data.features.reshape(self.data.features.shape[0], -1),
                  self.data.indices,
                  self.data.spatial_shape,
                  self.data.batch_size,
                  self.data.grid,
                  self.data.voxel_num,
                  self.data.indice_dict
              )
              new_data._features = feats
              new_data.benchmark = self.data.benchmark
              new_data.benchmark_record = self.data.benchmark_record
              new_data.thrust_allocator = self.data.thrust_allocator
              new_data._timer = self.data._timer
              new_data.force_algo = self.data.force_algo
              new_data.int8_scale = self.data.int8_scale
              if coords is not None:
                  new_data.indices = coords
          new_tensor = SparseTensor(new_data, shape=torch.Size(new_shape), layout=self.layout, scale=self._scale, spatial_cache=self._spatial_cache)
          return new_tensor
  
      @staticmethod
      def full(aabb, dim, value, dtype=torch.float32, device=None) -> 'SparseTensor':
          N, C = dim
          x = torch.arange(aabb[0], aabb[3] + 1)
          y = torch.arange(aabb[1], aabb[4] + 1)
          z = torch.arange(aabb[2], aabb[5] + 1)
          coords = torch.stack(torch.meshgrid(x, y, z, indexing='ij'), dim=-1).reshape(-1, 3)
          coords = torch.cat([
              torch.arange(N).view(-1, 1).repeat(1, coords.shape[0]).view(-1, 1),
              coords.repeat(N, 1),
          ], dim=1).to(dtype=torch.int32, device=device)
          feats = torch.full((coords.shape[0], C), value, dtype=dtype, device=device)
          return SparseTensor(feats=feats, coords=coords)
  
      def __merge_sparse_cache(self, other: 'SparseTensor') -> dict:
          new_cache = {}
          for k in set(list(self._spatial_cache.keys()) + list(other._spatial_cache.keys())):
              if k in self._spatial_cache:
                  new_cache[k] = self._spatial_cache[k]
              if k in other._spatial_cache:
                  if k not in new_cache:
                      new_cache[k] = other._spatial_cache[k]
                  else:
                      new_cache[k].update(other._spatial_cache[k])
          return new_cache
  
      def __neg__(self) -> 'SparseTensor':
          return self.replace(-self.feats)
      
      def __elemwise__(self, other: Union[torch.Tensor, 'SparseTensor'], op: callable) -> 'SparseTensor':
          if isinstance(other, torch.Tensor):
              try:
                  other = torch.broadcast_to(other, self.shape)
                  other = sparse_batch_broadcast(self, other)
              except:
                  pass
          if isinstance(other, SparseTensor):
              other = other.feats
          new_feats = op(self.feats, other)
          new_tensor = self.replace(new_feats)
          if isinstance(other, SparseTensor):
              new_tensor._spatial_cache = self.__merge_sparse_cache(other)
          return new_tensor
  
      def __add__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, torch.add)
  
      def __radd__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, torch.add)
      
      def __sub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, torch.sub)
      
      def __rsub__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, lambda x, y: torch.sub(y, x))
  
      def __mul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, torch.mul)
  
      def __rmul__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, torch.mul)
  
      def __truediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, torch.div)
  
      def __rtruediv__(self, other: Union[torch.Tensor, 'SparseTensor', float]) -> 'SparseTensor':
          return self.__elemwise__(other, lambda x, y: torch.div(y, x))
  
      def __getitem__(self, idx):
          if isinstance(idx, int):
              idx = [idx]
          elif isinstance(idx, slice):
              idx = range(*idx.indices(self.shape[0]))
          elif isinstance(idx, torch.Tensor):
              if idx.dtype == torch.bool:
                  assert idx.shape == (self.shape[0],), f"Invalid index shape: {idx.shape}"
                  idx = idx.nonzero().squeeze(1)
              elif idx.dtype in [torch.int32, torch.int64]:
                  assert len(idx.shape) == 1, f"Invalid index shape: {idx.shape}"
              else:
                  raise ValueError(f"Unknown index type: {idx.dtype}")
          else:
              raise ValueError(f"Unknown index type: {type(idx)}")
          
          coords = []
          feats = []
          for new_idx, old_idx in enumerate(idx):
              coords.append(self.coords[self.layout[old_idx]].clone())
              coords[-1][:, 0] = new_idx
              feats.append(self.feats[self.layout[old_idx]])
          coords = torch.cat(coords, dim=0).contiguous()
          feats = torch.cat(feats, dim=0).contiguous()
          return SparseTensor(feats=feats, coords=coords)
  
      def register_spatial_cache(self, key, value) -> None:
          """
          Register a spatial cache.
          The spatial cache can be any thing you want to cache.
          The registery and retrieval of the cache is based on current scale.
          """
          scale_key = str(self._scale)
          if scale_key not in self._spatial_cache:
              self._spatial_cache[scale_key] = {}
          self._spatial_cache[scale_key][key] = value
  
      def get_spatial_cache(self, key=None):
          """
          Get a spatial cache.
          """
          scale_key = str(self._scale)
          cur_scale_cache = self._spatial_cache.get(scale_key, {})
          if key is None:
              return cur_scale_cache
          return cur_scale_cache.get(key, None)
  
  
  def sparse_batch_broadcast(input: SparseTensor, other: torch.Tensor) -> torch.Tensor:
      """
      Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
      
      Args:
          input (torch.Tensor): 1D tensor to broadcast.
          target (SparseTensor): Sparse tensor to broadcast to.
          op (callable): Operation to perform after broadcasting. Defaults to torch.add.
      """
      coords, feats = input.coords, input.feats
      broadcasted = torch.zeros_like(feats)
      for k in range(input.shape[0]):
          broadcasted[input.layout[k]] = other[k]
      return broadcasted
  
  
  def sparse_batch_op(input: SparseTensor, other: torch.Tensor, op: callable = torch.add) -> SparseTensor:
      """
      Broadcast a 1D tensor to a sparse tensor along the batch dimension then perform an operation.
      
      Args:
          input (torch.Tensor): 1D tensor to broadcast.
          target (SparseTensor): Sparse tensor to broadcast to.
          op (callable): Operation to perform after broadcasting. Defaults to torch.add.
      """
      return input.replace(op(input.feats, sparse_batch_broadcast(input, other)))
  
  
  def sparse_cat(inputs: List[SparseTensor], dim: int = 0) -> SparseTensor:
      """
      Concatenate a list of sparse tensors.
      
      Args:
          inputs (List[SparseTensor]): List of sparse tensors to concatenate.
      """
      if dim == 0:
          start = 0
          coords = []
          for input in inputs:
              coords.append(input.coords.clone())
              coords[-1][:, 0] += start
              start += input.shape[0]
          coords = torch.cat(coords, dim=0)
          feats = torch.cat([input.feats for input in inputs], dim=0)
          output = SparseTensor(
              coords=coords,
              feats=feats,
          )
      else:
          feats = torch.cat([input.feats for input in inputs], dim=dim)
          output = inputs[0].replace(feats)
  
      return output
  
  
  def sparse_unbind(input: SparseTensor, dim: int) -> List[SparseTensor]:
      """
      Unbind a sparse tensor along a dimension.
      
      Args:
          input (SparseTensor): Sparse tensor to unbind.
          dim (int): Dimension to unbind.
      """
      if dim == 0:
          return [input[i] for i in range(input.shape[0])]
      else:
          feats = input.feats.unbind(dim)
          return [input.replace(f) for f in feats]
  