File size: 3,174 Bytes
ecc81b3 611aea1 ecc81b3 611aea1 ecc81b3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 | """torch-dimensions — N-dimensional sequence models for PyTorch.
Every model in scope is a 1-D mixer plus a plan for sweeping it over an N-D
lattice. See DESIGN.md for the architecture and PLAN.md for build order.
Conventionally imported as::
import torch_dimensions as td
Public names land here as each phase completes, and optional-dependency blocks
are imported defensively so a CPU-only install stays importable.
"""
from torch_dimensions import data, mixers, testing, viz
from torch_dimensions.compose import (
ND_METHODS,
AxialKernel,
AxialScan,
Flatten,
axial_apply,
axial_attention,
axial_contract,
axial_scan,
cafa,
flatten,
kron_operator,
register_nd_method,
resolve_nd_method,
)
from torch_dimensions.config import (
MODELS,
build,
list_models,
load,
read_config,
register_model,
save,
)
from torch_dimensions.lattice import Lattice, Restore, Sub
from torch_dimensions.mixers.conv import axis_receptive_field as receptive_field
from torch_dimensions.models.attention import Transformer, TransformerND
from torch_dimensions.models.conv import CNN, CNNND, TCN, TCNND
from torch_dimensions.models.rnn import GRU, LSTM
from torch_dimensions.models.ssm import (
S4,
S4D,
S4DND,
S4ND,
Mamba,
Mamba2,
Mamba2ND,
Mamba3,
Mamba3ND,
MambaND,
)
from torch_dimensions.models.vit import PatchEmbed, ViT
from torch_dimensions.optim import param_groups, warmup_cosine
from torch_dimensions.plan import AxisCoverage, Coverage, ScanPlan, Step
from torch_dimensions.spec import SPEC_VERSION as spec_version
from torch_dimensions.spec import spec
# Read from the installed metadata rather than written here. The literal that
# used to live at this line said 0.1.0 through the 0.2.0 and 0.3.1 releases —
# a published wheel that misreported its own version, which is exactly the
# thing a bug report quotes. One source of truth, in pyproject.toml.
try:
from importlib.metadata import PackageNotFoundError as _NotFound
from importlib.metadata import version as _dist_version
__version__ = _dist_version("torch-dimensions")
except _NotFound: # pragma: no cover - a source tree that was never installed
__version__ = "0.0.0+source"
__all__ = [
"CNN",
"CNNND",
"GRU",
"LSTM",
"TCN",
"TCNND",
"Mamba",
"Mamba2",
"Mamba2ND",
"Mamba3",
"Mamba3ND",
"MambaND",
"S4",
"S4D",
"S4DND",
"S4ND",
"PatchEmbed",
"param_groups",
"warmup_cosine",
"Transformer",
"TransformerND",
"ViT",
"MODELS",
"ND_METHODS",
"AxialKernel",
"AxialScan",
"Flatten",
"AxisCoverage",
"Coverage",
"Lattice",
"Restore",
"ScanPlan",
"Step",
"Sub",
"axial_apply",
"build",
"axial_attention",
"axial_contract",
"axial_scan",
"cafa",
"flatten",
"kron_operator",
"data",
"list_models",
"load",
"read_config",
"mixers",
"receptive_field",
"testing",
"viz",
"register_model",
"register_nd_method",
"save",
"resolve_nd_method",
"spec",
"spec_version",
]
|