Merge typing implementations, sort.

Signed-off-by: Emanuele Ballarin <[email protected]>
emaballarin · Dec 11, 2024 · 1f6df47 · 1f6df47
1 parent 9c5fcd4
commit 1f6df47
Show file tree

Hide file tree

Showing 11 changed files with 79 additions and 102 deletions.
diff --git a/ebtorch/__init__.py b/ebtorch/__init__.py
@@ -38,12 +38,12 @@
 del data_prep_dispatcher_3ch
 del fashionmnist_dataloader_dispatcher
 del imagenette_dataloader_dispatcher
-del mnist_dataloader_dispatcher
 del kmnist_dataloader_dispatcher
-del pathmnist_dataloader_dispatcher
+del mnist_dataloader_dispatcher
 del octmnist_dataloader_dispatcher
-del tissuemnist_dataloader_dispatcher
+del pathmnist_dataloader_dispatcher
 del tinyimagenet_dataloader_dispatcher
+del tissuemnist_dataloader_dispatcher
 
 # Deletions (from .distributed)
 del reduce_accumulate_keepalive
@@ -53,29 +53,31 @@
 del ArgMaxLayer
 del BasicAE
 del BasicVAE
-del SingleNeckVAE
-del beta_gaussian_kldiv
-del var_of_lap
 del BinarizeLayer
 del BrokenReLU
 del CausalConv1d
 del Clamp
+del Concatenate
+del ConvNeXtStem
+del ConvStem
 del ConvolutionalFlattenLayer
 del CoordConv1d
 del CoordConv2d
 del CoordConv3d
 del DeepRBL
+del DuplexLinearNeck
 del FCBlock
 del FieldTransform
 del FlatChannelize2DLayer
+del GRNorm
 del GaussianReparameterizerSampler
 del GaussianReparameterizerSamplerLegacy
+del GenerAct
 del InnerProduct
 del KWTA1d
 del KWTA2d
+del MetaAILayerNorm
 del Mish
-del multilasso
-del multiridge
 del MultiSolvePoissonTensor
 del NNEnsemble
 del PoissonNetCifar
@@ -87,47 +89,46 @@
 del SERLU
 del SGRUHCell
 del ScaledERF
+del SharedDuplexLinearNeck
+del SilhouetteScore
+del SimpleDuplexLinearNeck
 del SinLU
+del SingleNeckVAE
 del SirenSine
 del SmeLU
 del SolvePoisson
 del SolvePoissonTensor
+del StatefulTupleSelect
 del SwiGLU
 del TupleDecouple
 del TupleSelect
-del StatefulTupleSelect
-del SilhouetteScore
-del Concatenate
-del DuplexLinearNeck
-del SharedDuplexLinearNeck
-del SimpleDuplexLinearNeck
-del GenerAct
+del ViTStem
 del WideResNet
+del beta_gaussian_kldiv
 del beta_reco_bce
 del beta_reco_bce_splitout
-del lexsemble
+del beta_reco_mse
+del beta_reco_mse_splitout
+del bisided_thresholding
 del build_repeated_sequential
+del cummatmul
 del field_transform
+del lexsemble
+del logit_to_prob
 del mishlayer_init
+del multilasso
+del multiridge
+del oldtranspose
 del pixelwise_bce_mean
 del pixelwise_bce_sum
-del pixelwise_mse_sum
 del pixelwise_mse_mean
-del beta_reco_mse
-del beta_reco_mse_splitout
-del oldtranspose
+del pixelwise_mse_sum
 del silhouette_score
-del cummatmul
 del tensor_replicate
-del logit_to_prob
-del bisided_thresholding
-del ConvNeXtStem
-del ConvStem
-del GRNorm
-del MetaAILayerNorm
-del ViTStem
+del var_of_lap
 
 # Deletions (from .optim)
+del ADOPT
 del AdaBound
 del AdamP
 del Adan
@@ -137,19 +138,18 @@
 del Lookahead
 del Lookaround
 del RAdam
-del ADOPT
 del SAM
 del SGDP
+del alah_optim
 del epochwise_onecycle
+del make_beta_scheduler
 del onecycle_lincos
 del onecycle_linlin
 del onecycle_linlin_updown
-del warmed_up_linneal
-del warmed_up_annealer
-del make_beta_scheduler
-del alah_optim
 del ralah_optim
 del tricyc1c
+del warmed_up_annealer
+del warmed_up_linneal
 del wfneal
 
 # Deletions (from .logging)
@@ -159,4 +159,6 @@
 # Deletions (from .typing)
 del numlike
 del realnum
+del realnum
 del strdev
+del tensorlike
diff --git a/ebtorch/data/__init__.py b/ebtorch/data/__init__.py
@@ -31,4 +31,3 @@
 del cutmixup
 del datasets
 del prep
-del tinyimagenet
diff --git a/ebtorch/distributed/__init__.py b/ebtorch/distributed/__init__.py
@@ -15,5 +15,5 @@
 from .slurm import slurm_nccl_env
 
 # ──────────────────────────────────────────────────────────────────────────────
-del slurm
 del qol
+del slurm
diff --git a/ebtorch/nn/__init__.py b/ebtorch/nn/__init__.py
@@ -105,10 +105,10 @@
 del nnsemble
 del penalties
 del reshapelayers
+del serf
 del serlu
 del sinlu
 del smelu
-del serf
 
 # Deletions (from .functional)
 # del mish  # (already done by chance!)
@@ -118,49 +118,45 @@
 del AdverApply
 del AutoClipper
 del TA2ATAdapter
-del eval_model_on_test
-del extract_conv_filters
-del argser_f
-del find_lr
-del fxfx2module
-del act_opclone
-del suppress_std
 del TelegramBotEcho
-del stablediv
+del act_auto_broadcast
+del act_opclone
+del argser_f
 del argsink
-del no_op
+del broadcast_in_dim
+del custom_plot_setup
 del download_gdrive
 del emplace_kv
-del show_filters
-del subset_state_dict
+del eval_model_on_test
+del extract_conv_filters
+del find_lr
+del fromcache
+del fxfx2module
 del gather_model_repr
+del matched_apply
 del model_reqgrad
 del model_reqgrad_
+del no_op
 del patchify_2d
 del patchify_batch
 del patchify_dataset
+del petroff_2021_cmap
+del petroff_2021_color
+del petroff_2021_cycler
+del plot_out
+del repr_fx_flat_adapter
+del repr_sizes_flat_adapter
+del set_petroff_2021_colors
+del set_tableau10_colors
+del show_filters
+del stablediv
 del store_repr_autohook
 del store_repr_fx
 del store_repr_hook
-del petroff_2021_color
+del subset_state_dict
+del suppress_std
+del tableau10_cmap
 del tableau10_color
-del petroff_2021_cycler
 del tableau10_cycler
-del petroff_2021_cmap
-del tableau10_cmap
-del variadic_attrs
-del set_petroff_2021_colors
-del set_tableau10_colors
-del custom_plot_setup
-del plot_out
-del fromcache
-del repr_sizes_flat_adapter
-del repr_fx_flat_adapter
-del act_auto_broadcast
-del broadcast_in_dim
-del matched_apply
 del tensor_module_matched_apply
-del realnum
-del strdev
-del actvt
-del numlike
+del variadic_attrs
diff --git a/ebtorch/nn/utils/__init__.py b/ebtorch/nn/utils/__init__.py
@@ -19,10 +19,6 @@
 from .attrsutils import variadic_attrs
 from .autoclip import AutoClipper
 from .cacher import fromcache
-from .csttyping import actvt
-from .csttyping import numlike
-from .csttyping import realnum
-from .csttyping import strdev
 from .evalutils import eval_model_on_test
 from .filtermanip import extract_conv_filters
 from .filtermanip import show_filters
@@ -63,18 +59,17 @@
 from .reprutils import store_repr_hook
 
 # Deletions (from .)
-del adverutils
 del actab
+del adverutils
+del attrsutils
 del autoclip
+del cacher
 del evalutils
+del filtermanip
 del lrfinder
+del mapply
 del onlyutils
-del patches
-del reprutils
-del filtermanip
 del palettes
-del cacher
+del patches
 del plotting
-del attrsutils
-del mapply
-del csttyping
+del reprutils
diff --git a/ebtorch/nn/utils/csttyping.py b/ebtorch/nn/utils/csttyping.py
diff --git a/ebtorch/nn/utils/onlyutils.py b/ebtorch/nn/utils/onlyutils.py
@@ -41,8 +41,8 @@
 from torch import nn
 from torch import Tensor
 
-from .csttyping import actvt
-from .csttyping import numlike
+from ...typing import actvt
+from ...typing import numlike
 
 __all__ = [
     "argser_f",

diff --git a/ebtorch/optim/__init__.py b/ebtorch/optim/__init__.py
@@ -48,13 +48,13 @@
 # Deletions (from .)
 del adabound
 del adan
+del adopt
 del custom
 del lamb
 del lion
 del lookahead
 del lookaround
 del poptims
 del radam
-del adopt
 del sam
 del schedcos
diff --git a/ebtorch/typing/__init__.py b/ebtorch/typing/__init__.py
@@ -12,9 +12,11 @@
 #
 # ──────────────────────────────────────────────────────────────────────────────
 # Imports
+from .customtypes import actvt
 from .customtypes import numlike
 from .customtypes import realnum
 from .customtypes import strdev
+from .customtypes import tensorlike
 
 # Deletions (from .)
 del customtypes
diff --git a/ebtorch/typing/customtypes.py b/ebtorch/typing/customtypes.py
@@ -2,6 +2,7 @@
 # -*- coding: utf-8 -*-
 # ──────────────────────────────────────────────────────────────────────────────
 # Imports
+from collections.abc import Callable
 from typing import List
 from typing import Union
 
@@ -10,10 +11,11 @@
 from torch import Tensor
 
 # ──────────────────────────────────────────────────────────────────────────────
-__all__: List[str] = ["realnum", "strdev", "numlike", "tensorlike"]
+__all__: List[str] = ["realnum", "strdev", "numlike", "tensorlike", "actvt"]
 # ──────────────────────────────────────────────────────────────────────────────
 realnum = Union[int, float]
 strdev = Union[str, torch.device]
 tensorlike = Union[Tensor, np.ndarray]
 numlike = Union[realnum, tensorlike]
+actvt = Union[torch.nn.Module, Callable[[Tensor], Tensor]]
 # ──────────────────────────────────────────────────────────────────────────────
diff --git a/setup.py b/setup.py
@@ -26,7 +26,7 @@ def read(fname):
 
 setup(
     name=PACKAGENAME,
-    version="0.28.3",
+    version="0.28.4",
     author="Emanuele Ballarin",
     author_email="[email protected]",
     url="https://github.com/emaballarin/ebtorch",