torchsparsegradutils 0.1.0__tar.gz → 0.1.2__tar.gz

{torchsparsegradutils-0.1.0/torchsparsegradutils.egg-info → torchsparsegradutils-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchsparsegradutils
-Version: 0.1.0
+Version: 0.1.2
 Summary: A collection of utility functions to work with PyTorch sparse tensors
 Home-page: https://github.com/cai4cai/torchsparsegradutils
 Author: CAI4CAI research group
@@ -43,7 +43,7 @@ Things that are missing may be listed as [issues](https://github.com/cai4cai/tor
 ## Installation
 The provided package can be installed using:
 
-`pip install torchsparsegradutils` (TODO)
+`pip install torchsparsegradutils`
 
 or
 

{torchsparsegradutils-0.1.0 → torchsparsegradutils-0.1.2}/README.md

@@ -24,7 +24,7 @@ Things that are missing may be listed as [issues](https://github.com/cai4cai/tor
 ## Installation
 The provided package can be installed using:
 
-`pip install torchsparsegradutils` (TODO)
+`pip install torchsparsegradutils`
 
 or
 

{torchsparsegradutils-0.1.0 → torchsparsegradutils-0.1.2}/setup.py

@@ -8,7 +8,7 @@ def readme():
 
 setuptools.setup(
     name="torchsparsegradutils",
-    version="0.1.0",
+    version="0.1.2",
     description="A collection of utility functions to work with PyTorch sparse tensors",
     long_description=readme(),
     long_description_content_type="text/markdown",
@@ -19,7 +19,7 @@ setuptools.setup(
         "Programming Language :: Python :: 3.9",
         "Programming Language :: Python :: 3.10",
     ],
-    python_requires='>=3.8, <3.11',
+    python_requires=">=3.8, <3.11",
     keywords="sparse torch utility",
     url="https://github.com/cai4cai/torchsparsegradutils",
     author="CAI4CAI research group",
@@ -29,8 +29,8 @@ setuptools.setup(
     install_requires=[
         "torch>=1.13",
     ],
-    setup_requires=['pytest-runner'],
-    tests_require=['pytest'],
+    setup_requires=["pytest-runner"],
+    tests_require=["pytest"],
     test_suite="tests",
     extras_require={
         "extras": ["jax", "cupy"],

{torchsparsegradutils-0.1.0 → torchsparsegradutils-0.1.2}/torchsparsegradutils/distributions/sparse_multivariate_normal.py

@@ -76,6 +76,7 @@ class SparseMultivariateNormal(Distribution):
             in either torch.sparse_coo or torch.sparse_csr layout
     """
 
+    arg_constraints = {}
     # TODO: add in constraints
     # arg_constraints = {'loc': constraints.real_vector,
     #                    'diag': constraints.independent(constraints.positive, 1),

{torchsparsegradutils-0.1.0 → torchsparsegradutils-0.1.2}/torchsparsegradutils/encoders/pairwise_voxel_encoder.py

@@ -257,7 +257,7 @@ def calc_pariwise_coo_indices(
     return indices
 
 
-class PairwiseVoxelEncoder:
+class PairwiseVoxelEncoder(torch.nn.Module):
     """
     A class for encoding pairwise spatial local neighbourhoods and channel based voxel relations
     into a sparse tensor of either COO or CSR format.
@@ -270,6 +270,8 @@ class PairwiseVoxelEncoder:
     Additionally, diagonal entries can be included with the diag flag.
     The output matrix can be restricted to upper or lower triangular with the upper flag, if
     symmetric relationships are assumed, such as distance or correlation.
+    The indices are stored on the device specified by the device argument,
+    these indices can be sent to another device using the to() method.
 
     The sparse tensor is returned in the `__call__` method, which takes a tensor of values
     with shape [(B), N, C, H, D, W] and returns a sparse tensor of shape [(B), S, S]
@@ -284,6 +286,7 @@ class PairwiseVoxelEncoder:
     Some of the values will be trimmed from the 3D volume edges, as values are not allowed
     to wrap around the edges of the volume of spatial volume.
 
+    The sparse tensor is returned on the same device as the input values to the `__call__` method.
 
     Args:
         radius (float): The maximum distance from the origin within which the spatial
@@ -309,7 +312,8 @@ class PairwiseVoxelEncoder:
         indices_dtype (torch.dtype, optional): The data type of the output indices.
                                                Must be either torch.int32 or torch.int64.
                                                Default is torch.int64.
-        device (torch.device, optional): Device assigned to generate sparse tensor.
+        device (torch.device, optional): Device assigned to store sparse tensor indices
+                                         at initialisation.
                                          Defaults to torch.device("cpu").
 
 
@@ -333,6 +337,8 @@ class PairwiseVoxelEncoder:
         indices_dtype: torch.dtype = torch.int64,
         device: torch.device = torch.device("cpu"),
     ):
+        super().__init__()
+
         if not ((len(volume_shape) == 4) and all(isinstance(dim, int) and dim > 0 for dim in volume_shape)):
             raise ValueError("`volume_shape` must be a 4D tuple of positive integers, representing [C, H, D, W]")
 
@@ -348,7 +354,6 @@ class PairwiseVoxelEncoder:
         self.channel_voxel_relation = channel_voxel_relation
         self.layout = layout
         self.indices_dtype = indices_dtype
-        self.device = device
 
         self.volume_numel = reduce(mul, volume_shape)
 
@@ -369,6 +374,23 @@ class PairwiseVoxelEncoder:
         else:
             raise ValueError("layout must be either torch.sparse_coo or torch.sparse_csr")
 
+    def _apply(self, fn):
+        # Applying the function to the desired attributes
+        # This has been implemented to allow using the .to() method
+        for attr in ["indices", "csr_permutation", "crow_indices", "col_indices"]:
+            tensor = getattr(self, attr, None)
+            if tensor is not None:
+                setattr(self, attr, fn(tensor))
+
+        return self
+
+    @property
+    def device(self):
+        if self.layout == torch.sparse_coo:
+            return self.indices.device
+        elif self.layout == torch.sparse_csr:
+            return self.crow_indices.device
+
     def _calc_values(self, values: torch.Tensor) -> torch.Tensor:
         """
         Calculate the values for the sparse tensor based on the input values and offsets.
@@ -398,6 +420,8 @@ class PairwiseVoxelEncoder:
         boundaries of the input volume are trimmed to ensure that the output tensor doesn't
         contain any relationships that wrap around the spatial volume described by this encoder.
 
+        The output sparse tensor will be returned on the same device as the input values tensor.
+
         Args:
             values (torch.Tensor): Input tensor of values with shape [(B), N, C, H, D, W]
                                    where B is an optional batch dimension, N is the number of offsets
@@ -407,7 +431,7 @@ class PairwiseVoxelEncoder:
 
         Returns:
             torch.Tensor: Output tensor in either COO or CSR format, with shape [(B), S, S]
-            where is C*H*D*W.
+            where S is C*H*D*W.
 
         Raises:
             ValueError: If the shape of 'values' is not a 5D or 6D tensor.
@@ -452,7 +476,7 @@ class PairwiseVoxelEncoder:
             else:
                 sparse_dim_indices = self.indices.repeat(1, batch_size)
                 batch_dim_indices = (
-                    torch.arange(batch_size, dtype=self.indices.dtype, device=self.device)
+                    torch.arange(batch_size, dtype=self.indices.dtype, device=self.indices.device)
                     .repeat_interleave(self.indices.shape[-1])
                     .unsqueeze(0)
                 )

{torchsparsegradutils-0.1.0 → torchsparsegradutils-0.1.2}/torchsparsegradutils/tests/test_encoders.py

@@ -157,7 +157,7 @@ def test_trim(tensor_nd, offsets, expected_output_slice):
 current_dir = Path(os.path.abspath(os.path.dirname(__file__)))
 
 # Construct the path to the yaml file:
-yaml_file = current_dir / 'test_params' / 'xyz_coords.yaml'
+yaml_file = current_dir / "test_params" / "xyz_coords.yaml"
 
 with open(yaml_file) as f:  # load test cases from file
     coord_test_cases = yaml.safe_load(f)
@@ -179,7 +179,7 @@ def test_gen_coords(radius, expected_coords):
 
 # Test neighbourgood offset generation:
 
-yaml_file = current_dir / 'test_params' / 'czyx_shifts.yaml'
+yaml_file = current_dir / "test_params" / "czyx_shifts.yaml"
 
 with open(yaml_file) as f:  # load test cases from file
     shift_test_cases = yaml.safe_load(f)
@@ -261,7 +261,7 @@ def test_pairwise_coo_indices_unique(radius, volume_shape, diag, upper, channel_
 
 # Test the indices generated are as expected for a simple (3, 2, 2, 2) volume:
 
-yaml_file = current_dir / 'test_params' / 'pairwise_coo_indices.yaml'
+yaml_file = current_dir / "test_params" / "pairwise_coo_indices.yaml"
 
 with open(yaml_file) as f:  # load test cases from file to avoid a massive mess
     data = yaml.safe_load(f)
@@ -406,6 +406,25 @@ def test_PVE_dtype_device(batch_size, layout, indices_dtype, values_dtype, devic
         assert sparse_matrix.col_indices().device.type == device.type
 
 
+def test_PVE_to_device(layout, device):
+    encoder = PairwiseVoxelEncoder(
+        radius=1.0,
+        volume_shape=(5, 5, 5, 5),
+        layout=layout,
+        indices_dtype=torch.int64,
+        device=torch.device("cpu"),
+    )
+    encoder.to(device)
+
+    if layout == torch.sparse_coo:
+        assert encoder.indices.device.type == device.type
+
+    elif layout == torch.sparse_csr:
+        assert encoder.crow_indices.device.type == device.type
+        assert encoder.col_indices.device.type == device.type
+        assert encoder.csr_permutation.device.type == device.type
+
+
 # Test based on expected indices in pairwise_coo_indices.yaml
 @pytest.mark.parametrize(
     "radius, volume_shape, diag, upper, channel_relation, expected_indices",

{torchsparsegradutils-0.1.0 → torchsparsegradutils-0.1.2/torchsparsegradutils.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: torchsparsegradutils
-Version: 0.1.0
+Version: 0.1.2
 Summary: A collection of utility functions to work with PyTorch sparse tensors
 Home-page: https://github.com/cai4cai/torchsparsegradutils
 Author: CAI4CAI research group
@@ -43,7 +43,7 @@ Things that are missing may be listed as [issues](https://github.com/cai4cai/tor
 ## Installation
 The provided package can be installed using:
 
-`pip install torchsparsegradutils` (TODO)
+`pip install torchsparsegradutils`
 
 or