Type annotations for mlx.core module (#512)

This commit is contained in:
Danilo Peixoto 2024-01-21 17:53:12 -03:00 committed by GitHub
parent 7a34e46677
commit ecb174ca9d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 67 additions and 8 deletions

View File

@ -26,16 +26,21 @@ jobs:
command: | command: |
pip install --upgrade cmake pip install --upgrade cmake
pip install --upgrade pybind11[global] pip install --upgrade pybind11[global]
pip install pybind11-stubgen
pip install numpy pip install numpy
sudo apt-get update sudo apt-get update
sudo apt-get install libblas-dev sudo apt-get install libblas-dev
- run: - run:
name: Build python package name: Install Python package
command: | command: |
CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" CMAKE_BUILD_PARALLEL_LEVEL="" python3 setup.py build_ext --inplace CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" CMAKE_BUILD_PARALLEL_LEVEL="" python3 setup.py build_ext --inplace
CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" CMAKE_BUILD_PARALLEL_LEVEL="" python3 setup.py develop CMAKE_ARGS="-DMLX_BUILD_METAL=OFF" CMAKE_BUILD_PARALLEL_LEVEL="" python3 setup.py develop
- run: - run:
name: Run the python tests name: Generate package stubs
command: |
python3 setup.py generate_stubs
- run:
name: Run Python tests
command: | command: |
python3 -m unittest discover python/tests python3 -m unittest discover python/tests
# TODO: Reenable when extension api becomes stable # TODO: Reenable when extension api becomes stable
@ -65,19 +70,26 @@ jobs:
conda activate runner-env conda activate runner-env
pip install --upgrade cmake pip install --upgrade cmake
pip install --upgrade pybind11[global] pip install --upgrade pybind11[global]
pip install pybind11-stubgen
pip install numpy pip install numpy
pip install torch pip install torch
pip install tensorflow pip install tensorflow
pip install unittest-xml-reporting pip install unittest-xml-reporting
- run: - run:
name: Build python package name: Install Python package
command: | command: |
eval "$(conda shell.bash hook)" eval "$(conda shell.bash hook)"
conda activate runner-env conda activate runner-env
CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py build_ext --inplace CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py build_ext --inplace
CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py develop CMAKE_BUILD_PARALLEL_LEVEL="" python setup.py develop
- run: - run:
name: Run the python tests name: Generate package stubs
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
python setup.py generate_stubs
- run:
name: Run Python tests
command: | command: |
eval "$(conda shell.bash hook)" eval "$(conda shell.bash hook)"
conda activate runner-env conda activate runner-env
@ -121,10 +133,26 @@ jobs:
conda activate runner-env conda activate runner-env
pip install --upgrade cmake pip install --upgrade cmake
pip install --upgrade pybind11[global] pip install --upgrade pybind11[global]
pip install pybind11-stubgen
pip install numpy pip install numpy
pip install twine pip install twine
- run: - run:
name: Build package name: Install Python package
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
DEVELOPER_DIR=$(developer_dir_macos_<< parameters.macos_version >>) \
PYPI_RELEASE=1 \
CMAKE_BUILD_PARALLEL_LEVEL="" \
python setup.py install
- run:
name: Generate package stubs
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
python setup.py generate_stubs
- run:
name: Publish Python package
command: | command: |
eval "$(conda shell.bash hook)" eval "$(conda shell.bash hook)"
conda activate runner-env conda activate runner-env
@ -157,10 +185,26 @@ jobs:
conda activate runner-env conda activate runner-env
pip install --upgrade cmake pip install --upgrade cmake
pip install --upgrade pybind11[global] pip install --upgrade pybind11[global]
pip install pybind11-stubgen
pip install numpy pip install numpy
pip install twine pip install twine
- run: - run:
name: Build package name: Install Python package
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
DEVELOPER_DIR=$(developer_dir_macos_<< parameters.macos_version >>) \
DEV_RELEASE=1 \
CMAKE_BUILD_PARALLEL_LEVEL="" \
python setup.py install
- run:
name: Generate package stubs
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
python setup.py generate_stubs
- run:
name: Publish Python package
command: | command: |
eval "$(conda shell.bash hook)" eval "$(conda shell.bash hook)"
conda activate runner-env conda activate runner-env
@ -193,10 +237,25 @@ jobs:
conda activate runner-env conda activate runner-env
pip install --upgrade cmake pip install --upgrade cmake
pip install --upgrade pybind11[global] pip install --upgrade pybind11[global]
pip install pybind11-stubgen
pip install numpy pip install numpy
pip install twine pip install twine
- run: - run:
name: Build package name: Install Python package
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
DEVELOPER_DIR=$(developer_dir_macos_<< parameters.macos_version >>) \
CMAKE_BUILD_PARALLEL_LEVEL="" \
python setup.py install
- run:
name: Generate package stubs
command: |
eval "$(conda shell.bash hook)"
conda activate runner-env
python setup.py generate_stubs
- run:
name: Build package distribution
command: | command: |
eval "$(conda shell.bash hook)" eval "$(conda shell.bash hook)"
conda activate runner-env conda activate runner-env

View File

@ -148,7 +148,7 @@ if __name__ == "__main__":
where="python", exclude=["src", "tests", "tests.*"] where="python", exclude=["src", "tests", "tests.*"]
) )
package_dir = {"": "python"} package_dir = {"": "python"}
package_data = {"mlx": ["lib/*", "include/*", "share/*"]} package_data = {"mlx": ["lib/*", "include/*", "share/*"], "mlx.core": ["*.pyi"]}
setup( setup(
name="mlx", name="mlx",