flaxlib in cc

PiperOrigin-RevId: 696286840
google · Nov 13, 2024 · 5e45f71 · 5e45f71
1 parent ac3e85a
commit 5e45f71
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/flaxlib/flaxlib/__init__.py b/flaxlib/flaxlib/__init__.py
@@ -12,4 +12,4 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from flaxlib.flaxlib import sum_as_string as sum_as_string
+from flax.flaxlib.flaxlib_cc import sum_as_string as sum_as_string
diff --git a/flaxlib/flaxlib/flaxlib.pyi → flaxlib/flaxlib_cc.pyi b/flaxlib/flaxlib/flaxlib.pyi → flaxlib/flaxlib_cc.pyi
diff --git a/flaxlib/src/lib.cc b/flaxlib/src/lib.cc
@@ -0,0 +1,13 @@
+#include <string>
+
+#include "third_party/nanobind/include/nanobind/nanobind.h"
+
+namespace flaxlib {
+std::string sum_as_string(int a, int b) {
+  return std::to_string(a + b);
+}
+
+NB_MODULE(flaxlib_cc, m) {
+  m.def("sum_as_string", &sum_as_string);
+}
+}  // namespace flaxlib