Move import block

Author: iwyoo

chytorch/nn/losses.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Bool, Float
 #
 # Copyright 2023, 2024 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -25,7 +23,8 @@
 from torch import float32, zeros_like, exp, Tensor
 from torch.nn import Parameter, MSELoss
 from torch.nn.modules.loss import _Loss
-
+import torch
+from jaxtyping import Bool, Float
 
 
 class MultiTaskLoss(_Loss):

chytorch/nn/molecule/encoder.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Float
 #
 # Copyright 2021-2024 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -24,6 +22,8 @@
 #
 from itertools import repeat
 from torch.nn import GELU, Module, ModuleList, LayerNorm
+import torch
+from jaxtyping import Float
 
 from typing import Tuple, Optional, List
 from warnings import warn

chytorch/nn/reaction/encoder.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Float
 #
 # Copyright 2021-2023 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -25,6 +23,8 @@
 from math import inf
 from torch import zeros_like, float as t_float
 from torch.nn import Embedding, GELU, Module
+import torch
+from jaxtyping import Float
 
 from ..molecule import MoleculeEncoder
 from ..transformer import EncoderLayer

chytorch/nn/voting/binary.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Float, Int
 #
 # Copyright 2022, 2023 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -26,6 +24,8 @@
 from torch import sigmoid, no_grad
 from torch.nn import GELU
 from torch.nn.functional import binary_cross_entropy_with_logits
+import torch
+from jaxtyping import Float, Int
 
 from typing import Union, Optional
 from ._kfold import k_fold_mask

chytorch/nn/voting/classifier.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Float, Int
 #
 # Copyright 2022, 2023 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -26,6 +24,8 @@
 from torch import bmm, no_grad, Tensor
 from torch.nn import Dropout, GELU, LayerNorm, LazyLinear, Linear, Module
 from torch.nn.functional import cross_entropy, softmax
+import torch
+from jaxtyping import Float, Int
 
 from typing import Optional, Union
 from ._kfold import k_fold_mask

chytorch/nn/voting/regressor.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Float
 #
 # Copyright 2022, 2023 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -26,6 +24,8 @@
 from torch import bmm, no_grad, Tensor
 from torch.nn import Dropout, GELU, LayerNorm, LazyLinear, Linear, Module
 from torch.nn.functional import smooth_l1_loss
+import torch
+from jaxtyping import Float
 
 from typing import Optional, Union
 from ._kfold import k_fold_mask

chytorch/utils/data/molecule/conformer.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Int
 #
 # Copyright 2022-2024 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -29,6 +27,8 @@
 from torch import IntTensor, Size, zeros, ones as t_ones, int32 as t_int32, eye
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import Dataset
+import torch
+from jaxtyping import Int
 
 from typing import Sequence, Tuple, Union, NamedTuple
 from .._abc import default_collate_fn_map

chytorch/utils/data/molecule/encoder.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Int
 #
 # Copyright 2021-2024 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -28,6 +26,8 @@
 from torch import IntTensor, Size, int32, ones, zeros, eye, empty, full
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import Dataset
+import torch
+from jaxtyping import Int
 
 from typing import Sequence, Union, NamedTuple, Optional, Tuple
 from zlib import decompress

chytorch/utils/data/reaction/encoder.py

@@ -1,6 +1,4 @@
 # -*- coding: utf-8 -*-
-import torch
-from jaxtyping import Int
 #
 # Copyright 2021-2023 Ramil Nugmanov <nougmanoff@protonmail.com>
 #
@@ -27,6 +25,8 @@
 from torch import IntTensor, cat, zeros, int32, Size, eye
 from torch.nn.utils.rnn import pad_sequence
 from torch.utils.data import Dataset
+import torch
+from jaxtyping import Int
 
 from typing import Sequence, Union, NamedTuple
 from ..molecule import MoleculeDataset