Update GH Actions to support nanobind

openxla · Dec 12, 2024 · 082a659 · 082a659
1 parent 902cba7
commit 082a659
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/.github/workflows/buildAndTestCMake.yml b/.github/workflows/buildAndTestCMake.yml
@@ -88,7 +88,8 @@ jobs:
     - name: Build and Test StableHLO (with Python bindings)
       shell: bash
       run: |
-        pip install tensorflow-cpu
+        pip install tensorflow-cpu nanobind==2.4
+        pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
         ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
       env:
           CMAKE_BUILD_TYPE: Release

diff --git a/.github/workflows/publishWheelRelease.yml b/.github/workflows/publishWheelRelease.yml
@@ -83,7 +83,8 @@ jobs:
     - name: Build and Test StableHLO with Python
       shell: bash
       run: |
-        pip install tensorflow-cpu
+        pip install tensorflow-cpu nanobind==2.4
+        pip install -r "$LLVM_PROJECT_DIR/mlir/python/requirements.txt"
         ./build_tools/github_actions/ci_build_cmake.sh "$LLVM_BUILD_DIR" "$STABLEHLO_BUILD_DIR"
       env:
           CMAKE_BUILD_TYPE: Release

diff --git a/stablehlo/integrations/python/CMakeLists.txt b/stablehlo/integrations/python/CMakeLists.txt
@@ -94,6 +94,7 @@ declare_mlir_python_sources(CheckPythonExtensions)
 declare_mlir_python_extension(CheckPythonExtensions.Main
   MODULE_NAME _check
   ADD_TO_PARENT CheckPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
   CheckModule.cpp
   EMBED_CAPI_LINK_LIBS
@@ -106,6 +107,7 @@ declare_mlir_python_sources(ChloPythonExtensions)
 declare_mlir_python_extension(ChloPythonExtensions.Main
   MODULE_NAME _chlo
   ADD_TO_PARENT ChloPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
     ChloModule.cpp
   EMBED_CAPI_LINK_LIBS
@@ -118,6 +120,7 @@ declare_mlir_python_sources(StablehloPythonExtensions)
 declare_mlir_python_extension(StablehloPythonExtensions.Main
   MODULE_NAME _stablehlo
   ADD_TO_PARENT StablehloPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
     StablehloApi.cpp
     StablehloModule.cpp
@@ -132,6 +135,7 @@ declare_mlir_python_sources(VhloPythonExtensions)
 declare_mlir_python_extension(VhloPythonExtensions.Main
   MODULE_NAME _vhlo
   ADD_TO_PARENT VhloPythonExtensions
+  PYTHON_BINDINGS_LIBRARY nanobind
   SOURCES
     VhloModule.cpp
   EMBED_CAPI_LINK_LIBS

diff --git a/stablehlo/integrations/python/StablehloApi.cpp b/stablehlo/integrations/python/StablehloApi.cpp
@@ -259,7 +259,7 @@ void AddPortableApi(nb::module_ &m) {
       nb::arg("artifact_str"));
   m.def(
       "deserialize_portable_artifact_str",
-      [](const nb::bytes& artifact) -> nb::bytes {
+      [](const nb::bytes &artifact) -> nb::bytes {
         StringWriterHelper accumulator;
         if (mlirLogicalResultIsFailure(stablehloDeserializePortableArtifact(
                 toMlirStringRef(artifact), accumulator.getMlirStringCallback(),