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

package_dir = \
{'': 'src'}

packages = \
['x_mlps']

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

install_requires = \
['dm-haiku>=0.0.6,<0.0.7', 'einops>=0.4.1,<0.5.0']

setup_kwargs = {
    'name': 'x-mlps',
    'version': '0.5.0',
    'description': 'Configurable MLPs built on JAX and Haiku',
    'long_description': '# X-MLPs\n\nAn MLP model that provides a flexible foundation to implement, mix-and-match, and test various state-of-the-art MLP building blocks and architectures.\nBuilt on JAX and Haiku.\n\n## Installation\n\n```sh\npip install x-mlps\n```\n\n**Note**: X-MLPs will not install JAX for you (see [here](https://github.com/google/jax#installation) for install instructions).\n\n## Getting Started\n\nThe `XMLP` class provides the foundation from which all MLP architectures are built on, and is the primary class you use.\nAdditionally, X-MLPs relies heavily on factory functions to customize and instantiate the building blocks that make up a particular `XMLP` instance.\nFortunately, this library provides several SOTA MLP blocks out-of-the-box as factory functions.\nFor example, to implement the ResMLP architecture, you can implement the follow model function:\n\n```python\nimport haiku as hk\nimport jax\nfrom einops import rearrange\nfrom x_mlps import XMLP, Affine, resmlp_block_factory\n\ndef create_model(patch_size: int, dim: int, depth: int, num_classes: int = 10):\n    def model_fn(x: jnp.ndarray, is_training: bool) -> jnp.ndarray:\n        # Reformat input image into a sequence of patches\n        x = rearrange(x, "(h p1) (w p2) c -> (h w) (p1 p2 c)", p1=patch_size, p2=patch_size)\n        return XMLP(\n            num_patches=x.shape[-2],\n            dim=dim,\n            depth=depth,\n            block=resmlp_block_factory,\n            # Normalization following the stack of ResMLP blocks\n            normalization=lambda num_patches, dim, depth, **kwargs: Affine(dim, **kwargs),\n            num_classes=num_classes,\n        )(x, is_training=is_training)\n\n    # NOTE: Operating directly on batched data is supported as well.\n    return hk.vmap(model_fn, in_axes=(0, None))\n\nmodel = create_model(patch_size=4, dim=384, depth=12)\nmodel_fn = hk.transform(model)\n\nrng = jax.random.PRNGKey(0)\nparams = model_fn.init(rng, jnp.ones((1, 32, 32, 3)), False)\n```\n\nIt\'s important to note the `XMLP` module _does not_ reformat input data for you (e.g., to a sequence of patches).\nAs such, you must reformat data manually before feeding it to an `XMLP` module.\nThe [einops](https://github.com/arogozhnikov/einops) library, which is installed by X-MLPs, provides functions that can help here (e.g., `rearrange`).\n\n**Note**: Like the core Haiku modules, all modules implemented in X-MLPs support batched data and being vectorized via `vmap`.\n\n## X-MLPs Architecture Details\n\nX-MLPs uses a layered approach to construct arbitrary MLP networks. There are three core modules used to create a network\'s structure:\n\n1. `XSublayer` - bottom level module which wraps arbitrary feedforward functionality.\n2. `XBlock` - mid level module consisting of one or more `XSublayer` modules.\n3. `XMLP` - top level module which represents a generic MLP network, and is composed of a stack of repeated `XBlock` modules.\n\nTo support user-defined modules, each of the above modules support passing arbitrary keyword arguments to child modules.\nThis is accomplished by prepending arguments with one or more predefined prefixes (including user defined prefixes).\nBuilt-in prefixes include:\n\n1. "block\\_" - arguments fed directly to the `XBlock` module.\n2. "sublayers\\_" - arguments fed to all `XSublayer`s in each `XBlock`.\n3. "sublayers{i}\\_" - arguments fed to the i-th `XSublayer` in each `XBlock` (where 1 <= i <= # of sublayers).\n4. "ff\\_" - arguments fed to the feedforward module in a `XSublayer`.\n\nPrefixes must be combined in order when passing them to the `XMLP` module (e.g., "block_sublayer1_ff\\_<argument name>").\n\n### XSublayer\n\nThe `XSublayer` module is a flexible sublayer wrapper module providing skip connections and pre/post-normalization to an arbitrary child module (specifically, arbitrary feedforward modules e.g., `XChannelFeedForward`).\nChild module instances are not passed directly, rather a factory function which creates the child module is instead.\nThis ensures that individual sublayers can be configured automatically based on depth.\n\n### XBlock\n\nThe `XBlock` module is a generic MLP block. It is composed of one or more `XSublayer` modules, passed as factory functions.\n\n### XMLP\n\nAt the top level is the `XMLP` module, which represents a generic MLP network.\nN `XBlock` modules are stacked together to form a network, created via a common factory function.\n\n## Built-in MLP Architectures\n\nThe following architectures have been implemented in the form of `XBlock`s and have corresponding factory functions.\n\n- ResMLP - `resmlp_block_factory`\n- MLP-Mixer - `mlpmixer_block_factory`\n- gMLP - `gmlp_block_factory`\n- S²-MLP - `s2mlp_block_factory`\n- FNet - `fft_block_factory`\n\nImportantly, the components that make up these blocks are part of the public API, and thus can be easily mixed and\nmatched. Additionally, several component variations have been made based on combining ideas from current research.\nThis includes:\n\n- `fftlinear_block_factory` - an FNet block variation: combines a 1D FFT for patch mixing plus a linear layer.\n- `create_shift1d_op` - A 1D shift operation inspired by S²-MLP, making it appropriate for sequence data.\n- `create_multishift1d_op` - Like `create_shift1d_op`, but supports multiple shifts of varying sizes.\n\nSee their respective docstrings for more information.\n\n## LICENSE\n\nSee [LICENSE](LICENSE).\n\n## Citations\n\n```bibtex\n@article{Touvron2021ResMLPFN,\n  title={ResMLP: Feedforward networks for image classification with data-efficient training},\n  author={Hugo Touvron and Piotr Bojanowski and Mathilde Caron and Matthieu Cord and Alaaeldin El-Nouby and Edouard Grave and Gautier Izacard and Armand Joulin and Gabriel Synnaeve and Jakob Verbeek and Herv\'e J\'egou},\n  journal={ArXiv},\n  year={2021},\n  volume={abs/2105.03404}\n}\n```\n\n```bibtex\n@article{Tolstikhin2021MLPMixerAA,\n  title={MLP-Mixer: An all-MLP Architecture for Vision},\n  author={Ilya O. Tolstikhin and Neil Houlsby and Alexander Kolesnikov and Lucas Beyer and Xiaohua Zhai and Thomas Unterthiner and Jessica Yung and Daniel Keysers and Jakob Uszkoreit and Mario Lucic and Alexey Dosovitskiy},\n  journal={ArXiv},\n  year={2021},\n  volume={abs/2105.01601}\n}\n```\n\n```bibtex\n@article{Liu2021PayAT,\n  title={Pay Attention to MLPs},\n  author={Hanxiao Liu and Zihang Dai and David R. So and Quoc V. Le},\n  journal={ArXiv},\n  year={2021},\n  volume={abs/2105.08050}\n}\n```\n\n```bibtex\n@article{Yu2021S2MLPSM,\n  title={S2-MLP: Spatial-Shift MLP Architecture for Vision},\n  author={Tan Yu and Xu Li and Yunfeng Cai and Mingming Sun and Ping Li},\n  journal={ArXiv},\n  year={2021},\n  volume={abs/2106.07477}\n}\n```\n\n```bibtex\n@article{Touvron2021GoingDW,\n  title={Going deeper with Image Transformers},\n  author={Hugo Touvron and Matthieu Cord and Alexandre Sablayrolles and Gabriel Synnaeve and Herv\'e J\'egou},\n  journal={ArXiv},\n  year={2021},\n  volume={abs/2103.17239}\n}\n```\n\n```bibtex\n@inproceedings{Huang2016DeepNW,\n  title={Deep Networks with Stochastic Depth},\n  author={Gao Huang and Yu Sun and Zhuang Liu and Daniel Sedra and Kilian Q. Weinberger},\n  booktitle={ECCV},\n  year={2016}\n}\n```\n\n```bibtex\n@article{LeeThorp2021FNetMT,\n  title={FNet: Mixing Tokens with Fourier Transforms},\n  author={James P Lee-Thorp and Joshua Ainslie and Ilya Eckstein and Santiago Onta{\\~n}{\\\'o}n},\n  journal={ArXiv},\n  year={2021},\n  volume={abs/2105.03824}\n}\n```\n',
    'author': 'Miller Wilt',
    'author_email': 'miller@pyriteai.com',
    'maintainer': None,
    'maintainer_email': None,
    'url': 'https://github.com/PyriteAI/x-mlps',
    'package_dir': package_dir,
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.8,<4.0',
}


setup(**setup_kwargs)
