[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

perceiver_pytorch/__init__.py

@@ -1,3 +1,3 @@
-from perceiver_pytorch.perceiver_pytorch import Perceiver
-from perceiver_pytorch.perceiver_io import PerceiverIO
 from perceiver_pytorch.multi_perceiver_pytorch import MultiPerceiver
+from perceiver_pytorch.perceiver_io import PerceiverIO
+from perceiver_pytorch.perceiver_pytorch import Perceiver

perceiver_pytorch/convolutions.py

@@ -21,9 +21,7 @@ def __init__(
 
         layers = [self.make_layer(input_channels, output_channels, batch=use_batchnorm)]
         for _ in range(num_layers - 1):
-            layers += [
-                self.make_layer(output_channels, output_channels, batch=use_batchnorm)
-            ]
+            layers += [self.make_layer(output_channels, output_channels, batch=use_batchnorm)]
 
         super().__init__(*layers)
 
@@ -94,8 +92,7 @@ def __init__(
         # The decoder sets the number of upsamples as log2(upsample_value), and this changes the number of channels
         # in a similar way, so it all scales together.
         intermediate_output_channels = [
-            output_channels * pow(2, num_upsamples - 1 - i)
-            for i in range(0, num_upsamples)
+            output_channels * pow(2, num_upsamples - 1 - i) for i in range(0, num_upsamples)
         ]
         intermediate_input_channels = [input_channels] + intermediate_output_channels
 

perceiver_pytorch/decoders.py

@@ -1,7 +1,8 @@
-import torch
 import numpy as np
-from perceiver_pytorch.utils import reverse_space_to_depth
+import torch
+
 from perceiver_pytorch.convolutions import Conv2DUpsample, Conv3DUpsample
+from perceiver_pytorch.utils import reverse_space_to_depth
 
 
 class ImageDecoder(torch.nn.Module):
@@ -180,7 +181,9 @@ def forward(self, inputs: torch.Tensor) -> torch.Tensor:
 
 class ImageDecoderPatches(torch.nn.Module):
     def __init__(
-        self, spatial_upsample: int = 1, temporal_upsample: int = 1,
+        self,
+        spatial_upsample: int = 1,
+        temporal_upsample: int = 1,
     ):
         """
         Patch-based image decoder
@@ -196,7 +199,5 @@ def __init__(
         self.spatial_upsample = spatial_upsample
 
     def forward(self, inputs: torch.Tensor) -> torch.Tensor:
-        inputs = reverse_space_to_depth(
-            inputs, self.temporal_upsample, self.spatial_upsample
-        )
+        inputs = reverse_space_to_depth(inputs, self.temporal_upsample, self.spatial_upsample)
         return inputs

perceiver_pytorch/encoders.py

@@ -1,9 +1,10 @@
+import math
+
+import numpy as np
 import torch
-from torch import nn
-import torchvision
 import torch.nn.functional as F
-import numpy as np
-import math
+import torchvision
+from torch import nn
 
 from perceiver_pytorch.convolutions import Conv2DDownsample
 from perceiver_pytorch.utils import space_to_depth
@@ -66,9 +67,7 @@ def __init__(
                 spatial_downsample=spatial_downsample,
             )
         elif self.prep_type == "metnet":
-            self.encoder = ImageEncoderMetNet(
-                crop_size=crop_size, use_space2depth=use_space2depth
-            )
+            self.encoder = ImageEncoderMetNet(crop_size=crop_size, use_space2depth=use_space2depth)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         return self.encoder(x)
@@ -169,7 +168,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class ImageEncoderPatches(torch.nn.Module):
     def __init__(
-        self, spatial_downsample: int = 4, temporal_downsample: int = 1,
+        self,
+        spatial_downsample: int = 4,
+        temporal_downsample: int = 1,
     ):
         """
         Image encoder that uses patches
@@ -198,7 +199,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class ImageEncoderPixel(torch.nn.Module):
     def __init__(
-        self, spatial_downsample: int = 4, temporal_downsample: int = 1,
+        self,
+        spatial_downsample: int = 4,
+        temporal_downsample: int = 1,
     ):
         """
         Image encoder class for simple downsampling with pixels
@@ -231,7 +234,9 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
 
 class ImageEncoderMetNet(nn.Module):
     def __init__(
-        self, crop_size: int = 256, use_space2depth: bool = True,
+        self,
+        crop_size: int = 256,
+        use_space2depth: bool = True,
     ):
         """
         Performs the MetNet preprocessing of mean pooling Sat channels, followed by

perceiver_pytorch/gated.py

@@ -1,13 +1,11 @@
 import torch
-from torch import nn, einsum
 import torch.nn.functional as F
-
 from einops import rearrange, repeat
+from torch import einsum, nn
 
-from perceiver_pytorch.layers import exists, default, cache_fn, PreNorm, FeedForward, Attention
+from perceiver_pytorch.layers import Attention, FeedForward, PreNorm, cache_fn, default, exists
 from perceiver_pytorch.utils import fourier_encode
 
-
 # helpers
 
 
@@ -31,9 +29,7 @@ def forward(self, x, **kwargs):
         b, dim = x.shape[0], self.dim
         y = self.fn(x, **kwargs)
 
-        gated_output = self.gru(
-            rearrange(y, "... d -> (...) d"), rearrange(x, "... d -> (...) d")
-        )
+        gated_output = self.gru(rearrange(y, "... d -> (...) d"), rearrange(x, "... d -> (...) d"))
 
         gated_output = rearrange(gated_output, "(b n) d -> b n d", b=b)
         return gated_output
@@ -127,30 +123,22 @@ def __init__(
                 )
             )
 
-        self.to_logits = nn.Sequential(
-            nn.LayerNorm(latent_dim), nn.Linear(latent_dim, num_classes)
-        )
+        self.to_logits = nn.Sequential(nn.LayerNorm(latent_dim), nn.Linear(latent_dim, num_classes))
 
     def forward(self, data, mask=None):
         b, *axis, _, device = *data.shape, data.device
-        assert (
-            len(axis) == self.input_axis
-        ), "input data must have the right number of axis"
+        assert len(axis) == self.input_axis, "input data must have the right number of axis"
 
         # calculate fourier encoded positions in the range of [-1, 1], for all axis
 
         axis_pos = list(
             map(
-                lambda size: torch.linspace(
-                    -1.0, 1.0, steps=size, device=device
-                ),
+                lambda size: torch.linspace(-1.0, 1.0, steps=size, device=device),
                 axis,
             )
         )
         pos = torch.stack(torch.meshgrid(*axis_pos), dim=-1)
-        enc_pos = fourier_encode(
-            pos, self.max_freq, self.num_freq_bands, base=self.freq_base
-        )
+        enc_pos = fourier_encode(pos, self.max_freq, self.num_freq_bands, base=self.freq_base)
         enc_pos = rearrange(enc_pos, "... n d -> ... (n d)")
         enc_pos = repeat(enc_pos, "... -> b ...", b=b)
 

perceiver_pytorch/layers.py

@@ -2,7 +2,7 @@
 
 import torch
 from einops import rearrange, repeat
-from torch import nn, einsum
+from torch import einsum, nn
 from torch.nn import functional as F
 
 from perceiver_pytorch.rotary import apply_rotary_emb
@@ -37,9 +37,7 @@ def __init__(self, dim, fn, context_dim=None):
         super().__init__()
         self.fn = fn
         self.norm = nn.LayerNorm(dim)
-        self.norm_context = (
-            nn.LayerNorm(context_dim) if exists(context_dim) else None
-        )
+        self.norm_context = nn.LayerNorm(context_dim) if exists(context_dim) else None
 
     def forward(self, x, **kwargs):
         x = self.norm(x)
@@ -57,6 +55,7 @@ class GEGLU(nn.Module):
     Gaussian Error Gated Linear Unit.
     See Shazer 2020: https://arxiv.org/abs/2002.05202
     """
+
     def forward(self, x):
         x, gates = x.chunk(2, dim=-1)
         return x * F.gelu(gates)
@@ -85,9 +84,7 @@ def forward(self, x):
 
 
 class Attention(nn.Module):
-    def __init__(
-        self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0
-    ):
+    def __init__(self, query_dim, context_dim=None, heads=8, dim_head=64, dropout=0.0):
         """
         Args:
             query_dim: Size of the queries.
@@ -108,9 +105,7 @@ def __init__(
         self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
         self.to_kv = nn.Linear(context_dim, inner_dim * 2, bias=False)
 
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
-        )
+        self.to_out = nn.Sequential(nn.Linear(inner_dim, query_dim), nn.Dropout(dropout))
 
     def forward(self, x, context=None, mask=None, pos_emb=None):
         """
@@ -134,9 +129,7 @@ def forward(self, x, context=None, mask=None, pos_emb=None):
         # Rearrange the query, key and value tensors.
         # b = batch size; n = TODO (PD-2021-09-13)
         # h = number of heads; d = number of dims per head.
-        q, k, v = map(
-            lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v)
-        )
+        q, k, v = map(lambda t: rearrange(t, "b n (h d) -> (b h) n d", h=h), (q, k, v))
 
         if exists(pos_emb):
             q, k = apply_rotary_emb(q, k, pos_emb)

perceiver_pytorch/mixed_latents.py

@@ -1,13 +1,11 @@
 import torch
-from torch import nn, einsum
 import torch.nn.functional as F
-
 from einops import rearrange, repeat
+from torch import einsum, nn
 
-from perceiver_pytorch.layers import exists, default, cache_fn, PreNorm, FeedForward, Attention
+from perceiver_pytorch.layers import Attention, FeedForward, PreNorm, cache_fn, default, exists
 from perceiver_pytorch.utils import fourier_encode
 
-
 # latent mixer
 
 

perceiver_pytorch/modalities.py

@@ -1,6 +1,7 @@
-import torch
 from dataclasses import dataclass
 
+import torch
+
 
 @dataclass
 class InputModality:

perceiver_pytorch/multi_perceiver_pytorch.py

@@ -1,10 +1,12 @@
-from perceiver_pytorch.perceiver_io import PerceiverIO
-from perceiver_pytorch.modalities import InputModality, modality_encoding
-from perceiver_pytorch.utils import encode_position, fourier_encode
+from math import prod
+from typing import Any, Dict, Iterable, List, Optional, Tuple, Union
+
 import torch
-from typing import List, Iterable, Dict, Optional, Any, Union, Tuple
 from einops import rearrange, repeat
-from math import prod
+
+from perceiver_pytorch.modalities import InputModality, modality_encoding
+from perceiver_pytorch.perceiver_io import PerceiverIO
+from perceiver_pytorch.utils import encode_position, fourier_encode
 
 
 class MultiPerceiver(torch.nn.Module):