Commit 219dfbf9 authored by Johan Gudmundsson's avatar Johan Gudmundsson
Browse files

update __all__ in torch_proxies

parent c792f688

src/borch/nn/torch_proxies.py

+48 −2
Original line number Diff line number Diff line
@@ -240,6 +240,7 @@ __all__ = [ 
    "AdaptiveAvgPool1d",

    "AdaptiveAvgPool2d",

    "AdaptiveAvgPool3d",

    "AdaptiveLogSoftmaxWithLoss",

    "AdaptiveMaxPool1d",

    "AdaptiveMaxPool2d",

    "AdaptiveMaxPool3d",
@@ -253,9 +254,13 @@ __all__ = [ 
    "BatchNorm2d",

    "BatchNorm3d",

    "Bilinear",

    "CELU",

    "CTCLoss",

    "ChannelShuffle",

    "ConstantPad1d",

    "ConstantPad2d",

    "ConstantPad3d",

    "Container",

    "Conv1d",

    "Conv2d",

    "Conv3d",
@@ -273,16 +278,23 @@ __all__ = [ 
    "ELU",

    "Embedding",

    "EmbeddingBag",

    "Fold",

    "FeatureAlphaDropout",

    "Flatten",

    "Fold",

    "FractionalMaxPool2d",

    "FractionalMaxPool3d",

    "GELU",

    "GLU",

    "GRU",

    "GRUCell",

    "GaussianNLLLoss",

    "GroupNorm",

    "Hardshrink",

    "Hardsigmoid",

    "Hardswish",

    "Hardtanh",

    "HingeEmbeddingLoss",

    "Identity",

    "InstanceNorm1d",

    "InstanceNorm2d",

    "InstanceNorm3d",
@@ -293,6 +305,13 @@ __all__ = [ 
    "LSTM",

    "LSTMCell",

    "LayerNorm",

    "LazyConv1d",

    "LazyConv2d",

    "LazyConv3d",

    "LazyConvTranspose1d",

    "LazyConvTranspose2d",

    "LazyConvTranspose3d",

    "LazyLinear",

    "LeakyReLU",

    "Linear",

    "LocalResponseNorm",
@@ -311,16 +330,20 @@ __all__ = [ 
    "MultiLabelMarginLoss",

    "MultiLabelSoftMarginLoss",

    "MultiMarginLoss",

    "MultiheadAttention",

    "NLLLoss",

    "NLLLoss2d",

    "PReLU",

    "PairwiseDistance",

    "ParameterDict",

    "ParameterList",

    "PixelShuffle",

    "PixelUnshuffle",

    "PoissonNLLLoss",

    "RNN",

    "RNNBase",

    "RNNCell",

    "RNNCellBase",

    "RReLU",

    "ReLU",

    "ReLU6",
@@ -330,8 +353,9 @@ __all__ = [ 
    "ReplicationPad2d",

    "ReplicationPad3d",

    "SELU",

    "Sigmoid",

    "Sequential",

    "SiLU",

    "Sigmoid",

    "SmoothL1Loss",

    "SoftMarginLoss",

    "Softmax",
@@ -340,13 +364,35 @@ __all__ = [ 
    "Softplus",

    "Softshrink",

    "Softsign",

    "SyncBatchNorm",

    "Tanh",

    "Tanhshrink",

    "Threshold",

    "TransformerDecoder",

    "TransformerDecoderLayer",

    "TransformerEncoder",

    "TransformerEncoderLayer",

    "TripletMarginLoss",

    "TripletMarginWithDistanceLoss",

    "Unflatten",

    "Unfold",

    "Upsample",

    "UpsamplingBilinear2d",

    "UpsamplingNearest2d",

    "ZeroPad2d",

]



_NEWER_MODULES = {

    "HuberLoss",

    "LazyBatchNorm1d",

    "LazyBatchNorm2d",

    "LazyBatchNorm3d",

    "LazyInstanceNorm1d",

    "LazyInstanceNorm2d",

    "LazyInstanceNorm3d",

    "Mish",

    "ReflectionPad3d",

}

# For backwards compatibility we can not have _NEWER_MODULES in __all__ as they might

# not exists in order versions of torch

__all__.extend(list(_NEWER_MODULES.intersection(set(globals()))))