# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['torchmix',
 'torchmix.components',
 'torchmix.components.layers',
 'torchmix.core',
 'torchmix.nn',
 'torchmix.third_party']

package_data = \
{'': ['*']}

install_requires = \
['einops>=0.6.0,<0.7.0',
 'hydra-core>=1.3.1,<2.0.0',
 'hydra-zen>=0.8.0,<0.9.0',
 'jaxtyping>=0.2.11,<0.3.0']

setup_kwargs = {
    'name': 'torchmix',
    'version': '0.1.0rc3',
    'description': 'A collection of useful PyTorch modules',
    'long_description': '<h1 align="center">torchmix</h1>\n\n<h3 align="center">The missing component library for PyTorch</h3>\n\n<br />\n\nWelcome to `torchmix`, a collection of PyTorch modules that aims to make your code more efficient and modular. We\'ve included a range of operations, from basic ones like `Repeat` and `Add`, to more complex ones like `WindowAttention` in the [Swin-Transformer](https://arxiv.org/abs/2103.14030). Our goal is to make it easy for you to use these various operations with minimal code, so you can focus on building your project rather than writing boilerplate.\n\nWe\'ve designed `torchmix` to be as user-friendly as possible. Each implementation is kept minimal and easy to understand, using [`einops`](https://github.com/arogozhnikov/einops) to avoid confusing tensor manipulation (such as `permute`, `transpose`, and `reshape`) and [`jaxtyping`](https://github.com/google/jaxtyping) to clearly document the shapes of the input and output tensors. This means that you can use `torchmix` with confidence, knowing that the components you\'re working with are clean and reliable.\n\n**Note: `torchmix` is a prototype that is currently in development and has not been tested for production use. The API may change at any time.**\n\n<br />\n\n## Install\n\nTo use `torchmix`, you will need to have `torch` already installed on your environment.\n\n```sh\npip install torchmix\n```\n\n## Usage\n\nTo use `torchmix`, simply import the desired components:\n\n```python\nimport torchmix.nn as nn  # Wrapped version of torch.nn\nfrom torchmix import (\n    Add,\n    Attach,\n    AvgPool,\n    ChannelMixer,\n    Extract,\n    PatchEmbed,\n    PositionEmbed,\n    PreNorm,\n    Repeat,\n    SelfAttention,\n    Token,\n)\n```\n\nYou can simply compose this components to build more complex architecture:\n\n```python\n# ViT with CLS Token attached\nvit_cls = nn.Sequential(\n    Add(\n        Attach(\n            PatchEmbed(dim=1024),\n            Token(dim=1024),\n        ),\n        PositionEmbed(\n            seq_length=196 + 1,\n            dim=1024,\n        ),\n    ),\n    Repeat(\n        nn.Sequential(\n            PreNorm(\n                ChannelMixer(\n                    dim=1024,\n                    expansion_factor=4,\n                    act_layer=nn.GELU.partial(),\n                ),\n                dim=1024,\n            ),\n            PreNorm(\n                SelfAttention(\n                    dim=1024,\n                    num_heads=8,\n                    head_dim=64,\n                ),\n                dim=1024,\n            ),\n        ),\n        depth=24,\n    ),\n    Extract(0),\n)\n\n# ViT with average pooling\nvit_avg = nn.Sequential(\n    Add(\n        PatchEmbed(dim=1024),\n        PositionEmbed(\n            seq_length=196,\n            dim=1024,\n        ),\n    ),\n    Repeat(\n        nn.Sequential(\n            PreNorm(\n                ChannelMixer(\n                    dim=1024,\n                    expansion_factor=4,\n                    act_layer=nn.GELU.partial(),\n                ),\n                dim=1024,\n            ),\n            PreNorm(\n                SelfAttention(\n                    dim=1024,\n                    num_heads=8,\n                    head_dim=64,\n                ),\n                dim=1024,\n            ),\n        ),\n        depth=24,\n    ),\n    AvgPool(),\n)\n```\n\n### Integration with [Hydra](https://hydra.cc/)\n\nReproducibility is important, so it is always a good idea to manage the configurations of your models. However, manually writing the configurations for complex, deeply nested PyTorch modules can be time-consuming and result in confusing and hard-to-maintain code. This is often because the parent class is responsible for accepting and passing along the parameters to its children classes, leading to a large number of arguments and strong coupling between the parent and children classes.\n\n`torchmix` was designed to address this issue by [**auto-magically**](https://github.com/mit-ll-responsible-ai/hydra-zen) generating the full configuration of a PyTorch modules **simply by instantiating them**, regardless of how deeply they are nested. This makes it easy to integrate your favorite modules into [`hydra`](https://hydra.cc/) ecosystem. This **instantiate-for-configuration** pattern also promotes the direct injection of dependencies, leading to the creation of loosely-coupled components and more declarative and intuitive code.\n\nIn other words, getting a configuration with `torchmix` is practically effortless - it\'s just there for you:\n\n```python\nfrom torchmix import nn\n\nmodel = nn.Sequential(\n    nn.Linear(1024, 4096),\n    nn.Dropout(0.1),\n    nn.GELU(),\n    nn.Linear(4096, 1024),\n    nn.Dropout(0.1),\n)\n\nmodel.config  # DictConfig which contains full signatures 🤯\n```\n\nYou can then store the configuration in the [`hydra`\'s `ConfigStore`](https://hydra.cc/docs/tutorials/structured_config/config_store/) using:\n\n```python\nmodel.store(group="model", name="mlp")\n```\n\nAlternatively, you can export it to a YAML file if you want:\n\n```python\nmodel.export("mlp.yaml")\n```\n\nThis will generate the following configuration:\n\n```yaml\n_target_: torchmix.nn.Sequential\n_args_:\n  - _target_: torchmix.nn.Linear\n    in_features: 1024\n    out_features: 4096\n    bias: true\n    device: null\n    dtype: null\n  - _target_: torchmix.nn.Dropout\n    p: 0.1\n    inplace: false\n  - _target_: torchmix.nn.GELU\n    approximate: none\n  - _target_: torchmix.nn.Linear\n    in_features: 4096\n    out_features: 1024\n    bias: true\n    device: null\n    dtype: null\n  - _target_: torchmix.nn.Dropout\n    p: 0.1\n    inplace: false\n```\n\nYou can always get the actual PyTorch module from its configuration using [`hydra`\'s `instantiate`.](https://hydra.cc/docs/advanced/instantiate_objects/overview/)\n\nTo create custom modules with this functionality, simply subclass `MixModule` and define your module as you normally would:\n\n```python\nfrom torchmix import MixModule\n\nclass CustomModule(MixModule):\n    def __init__(self, num_heads, dim, depth):\n        pass\n\ncustom_module = CustomModule(16, 768, 12)\ncustom_module.store(group="model", name="custom")\n```\n\n## Examples\n\nFor more information on using `torchmix`\'s functionality, see the following examples:\n\n- [Using the `no_parameters` Context Manager](https://github.com/torchmix/torchmix/tree/main/examples/no_parameters)\n- [Using the `partial` Method](https://github.com/torchmix/torchmix/tree/main/examples/partial)\n\n## Documentation\n\nDocumentation is currently in progress. Please stay tuned! 🚀\n\n## Contributing\n\nThe development of `torchmix` is an open process, and we welcome any contributions or suggestions for improvement. If you have ideas for new components or ways to enhance the library, feel free to open an issue or start a discussion. We welcome all forms of feedback, including criticism and suggestions for significant design changes. Please note that `torchmix` is currently in the early stages of development and any contributions should be considered experimental. Thank you for your support of `torchmix`!\n\n## License\n\n`torchmix` is licensed under the MIT License.\n',
    'author': 'junhsss',
    'author_email': 'junhsssr@gmail.com',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'None',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.8,<4.0',
}


setup(**setup_kwargs)
