mirror of
https://github.com/MLSysBook/TinyTorch.git
synced 2026-04-28 07:27:39 -05:00
c955437078
Stage 5 of TinyTorch API simplification: - Created tinytorch.nn package with PyTorch-compatible interface - Added Module base class in nn.modules for automatic parameter registration - Added functional module with relu, flatten, max_pool2d operations - Created tinytorch.optim package exposing Adam and SGD optimizers - Updated main __init__.py to export nn and optim modules - Linear and Conv2d now available through clean nn interface Students can now write PyTorch-like code: import tinytorch.nn as nn import tinytorch.nn.functional as F model = nn.Linear(784, 10) x = F.relu(model(x))
22 lines
346 B
Python
Generated
22 lines
346 B
Python
Generated
__version__ = "0.1.0"
|
|
|
|
# Import core functionality
|
|
from . import core
|
|
|
|
# Import PyTorch-compatible modules
|
|
from . import nn
|
|
from . import optim
|
|
|
|
# Make common components easily accessible
|
|
from .core.tensor import Tensor
|
|
from .nn import Module
|
|
|
|
# Export main public API
|
|
__all__ = [
|
|
'core',
|
|
'nn',
|
|
'optim',
|
|
'Tensor',
|
|
'Module'
|
|
]
|