flaxlib in cc

PiperOrigin-RevId: 696286840
google · Nov 14, 2024 · 0b19b33 · 0b19b33
1 parent ac3e85a
commit 0b19b33
Show file tree

Hide file tree

Showing 15 changed files with 37 additions and 24 deletions.
diff --git a/.github/workflows/flax_test.yml b/.github/workflows/flax_test.yml
@@ -108,13 +108,10 @@ jobs:
       uses: astral-sh/setup-uv@v2
       with:
         version: "0.3.0"
-    - name: Setup Rust (flaxlib)
-      uses: actions-rust-lang/setup-rust-toolchain@v1
 
     - name: Install dependencies
       run: |
         uv sync --extra all --extra testing --extra docs
-        uv pip install ./flaxlib
     - name: Install JAX
       run: |
         if [[ "${{ matrix.jax-version }}" == "newest" ]]; then

diff --git a/.gitignore b/.gitignore
@@ -13,6 +13,10 @@ build/
 docs*/**/_autosummary
 docs*/_build
 docs*/**/tmp
+flaxlib_src/build
+flaxlib_src/builddir
+flaxlib_src/dist
+flaxlib_src/subprojects
 
 # used by direnv
 .envrc

diff --git a/flaxlib/flaxlib/__init__.py b/flaxlib/flaxlib/__init__.py
diff --git a/flaxlib/.gitignore → flaxlib_src/.gitignore b/flaxlib/.gitignore → flaxlib_src/.gitignore
diff --git a/flaxlib/Cargo.lock → flaxlib_src/Cargo.lock b/flaxlib/Cargo.lock → flaxlib_src/Cargo.lock
diff --git a/flaxlib/Cargo.toml → flaxlib_src/Cargo.toml b/flaxlib/Cargo.toml → flaxlib_src/Cargo.toml
diff --git a/flaxlib/LICENSE → flaxlib_src/LICENSE b/flaxlib/LICENSE → flaxlib_src/LICENSE
diff --git a/flaxlib/README.md → flaxlib_src/README.md b/flaxlib/README.md → flaxlib_src/README.md
diff --git a/flaxlib/flaxlib/flaxlib.pyi → flaxlib_src/flaxlib.pyi b/flaxlib/flaxlib/flaxlib.pyi → flaxlib_src/flaxlib.pyi
diff --git a/flaxlib_src/meson.build b/flaxlib_src/meson.build
@@ -0,0 +1,14 @@
+project(
+  'flaxlib',
+  'cpp',
+  version: '0.0.1',
+  default_options: ['cpp_std=c++17'],
+)
+py = import('python').find_installation()
+nanobind_dep = dependency('nanobind', static: true)
+py.extension_module(
+  'flaxlib',
+  sources: ['src/lib.cc'],
+  dependencies: [nanobind_dep],
+  install: true,
+)
diff --git a/flaxlib/pyproject.toml → flaxlib_src/pyproject.toml b/flaxlib/pyproject.toml → flaxlib_src/pyproject.toml
@@ -1,12 +1,12 @@
 [build-system]
-requires = ["maturin>=1.7,<2.0"]
-build-backend = "maturin"
+requires = ['meson-python']
+build-backend = 'mesonpy'
 
 [project]
 name = "flaxlib"
 requires-python = ">=3.10"
 classifiers = [
-    "Programming Language :: Rust",
+    "Programming Language :: C++",
     "Programming Language :: Python :: Implementation :: CPython",
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
@@ -15,5 +15,3 @@ dynamic = ["version"]
 tests = [
     "pytest",
 ]
-[tool.maturin]
-features = ["pyo3/extension-module"]
diff --git a/flaxlib_src/src/lib.cc b/flaxlib_src/src/lib.cc
@@ -0,0 +1,14 @@
+#include <string>
+
+#include "nanobind/nanobind.h"
+#include "nanobind/stl/string.h"
+
+namespace flaxlib {
+std::string sum_as_string(int a, int b) {
+  return std::to_string(a + b);
+}
+
+NB_MODULE(flaxlib, m) {
+  m.def("sum_as_string", &sum_as_string);
+}
+}  // namespace flaxlib
diff --git a/flaxlib/src/lib.rs → flaxlib_src/src/lib.rs b/flaxlib/src/lib.rs → flaxlib_src/src/lib.rs
diff --git a/flaxlib/uv.lock → flaxlib_src/uv.lock b/flaxlib/uv.lock → flaxlib_src/uv.lock
diff --git a/tests/run_all_tests.sh b/tests/run_all_tests.sh
@@ -113,7 +113,8 @@ if $RUN_PYTEST; then
   echo "=== RUNNING PYTESTS ==="
   # Run some test on separate process, avoiding device configs poluting each other
   PYTEST_IGNORE=
-  for file in "tests/jax_utils_test.py"; do
+  # TODO(ivyzheng): Remove flaxlib_test.py once we get CI build for flaxlib CC.
+  for file in "tests/jax_utils_test.py" "tests/flaxlib_test.py"; do
       echo "pytest -n auto $file $PYTEST_OPTS"
       pytest -n auto $file $PYTEST_OPTS
       PYTEST_IGNORE+=" --ignore=$file"