vector-quantize-pytorch 1.16.2__tar.gz → 1.17.0__tar.gz

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: vector-quantize-pytorch
-Version: 1.16.2
+Version: 1.17.0
 Summary: Vector Quantization - Pytorch
 Project-URL: Homepage, https://pypi.org/project/vector-quantize-pytorch/
 Project-URL: Repository, https://github.com/lucidrains/vector-quantizer-pytorch

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "vector-quantize-pytorch"
-version = "1.16.2"
+version = "1.17.0"
 description = "Vector Quantization - Pytorch"
 authors = [
     { name = "Phil Wang", email = "lucidrains@gmail.com" }

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/vector_quantize_pytorch/finite_scalar_quantization.py

@@ -12,7 +12,7 @@ import torch
 import torch.nn as nn
 from torch.nn import Module
 from torch import Tensor, int32
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from einops import rearrange, pack, unpack
 
@@ -159,7 +159,7 @@ class FSQ(Module):
 
         return codes
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def forward(self, z):
         """
         einstein notation

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/vector_quantize_pytorch/lookup_free_quantization.py

@@ -16,7 +16,7 @@ import torch
 from torch import nn, einsum
 import torch.nn.functional as F
 from torch.nn import Module
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from einops import rearrange, reduce, pack, unpack
 
@@ -36,7 +36,7 @@ def maybe_distributed_mean(t):
     if not is_distributed():
         return t
 
-    dist.all_reduce(t)
+    dist.nn.all_reduce(t)
     t = t / dist.get_world_size()
     return t
 
@@ -293,7 +293,7 @@ class LFQ(Module):
 
         force_f32 = self.force_quantization_f32
 
-        quantization_context = partial(autocast, enabled = False) if force_f32 else nullcontext
+        quantization_context = partial(autocast, 'cuda', enabled = False) if force_f32 else nullcontext
 
         with quantization_context():
 

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/vector_quantize_pytorch/residual_fsq.py

@@ -8,7 +8,7 @@ import torch
 from torch import nn
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from vector_quantize_pytorch.finite_scalar_quantization import FSQ
 
@@ -167,7 +167,7 @@ class ResidualFSQ(Module):
 
         # go through the layers
 
-        with autocast(enabled = False):
+        with autocast('cuda', enabled = False):
             for quantizer_index, (layer, scale) in enumerate(zip(self.layers, self.scales)):
 
                 if should_quantize_dropout and quantizer_index > rand_quantize_dropout_index:

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/vector_quantize_pytorch/residual_lfq.py

@@ -6,7 +6,7 @@ import torch
 from torch import nn
 from torch.nn import Module, ModuleList
 import torch.nn.functional as F
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 from vector_quantize_pytorch.lookup_free_quantization import LFQ
 
@@ -156,7 +156,7 @@ class ResidualLFQ(Module):
 
         # go through the layers
 
-        with autocast(enabled = False):
+        with autocast('cuda', enabled = False):
             for quantizer_index, layer in enumerate(self.layers):
 
                 if should_quantize_dropout and quantizer_index > rand_quantize_dropout_index:

{vector_quantize_pytorch-1.16.2 → vector_quantize_pytorch-1.17.0}/vector_quantize_pytorch/vector_quantize_pytorch.py

@@ -9,7 +9,7 @@ from torch import nn, einsum, Tensor
 import torch.nn.functional as F
 import torch.distributed as distributed
 from torch.optim import Optimizer
-from torch.cuda.amp import autocast
+from torch.amp import autocast
 
 import einx
 from einops import rearrange, repeat, reduce, pack, unpack
@@ -458,7 +458,7 @@ class EuclideanCodebook(Module):
         batch_samples = rearrange(batch_samples, 'h ... d -> h (...) d')
         self.replace(batch_samples, batch_mask = expired_codes)
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def forward(
         self,
         x,
@@ -671,7 +671,7 @@ class CosineSimCodebook(Module):
         batch_samples = rearrange(batch_samples, 'h ... d -> h (...) d')
         self.replace(batch_samples, batch_mask = expired_codes)
 
-    @autocast(enabled = False)
+    @autocast('cuda', enabled = False)
     def forward(
         self,
         x,