flaxlib in cc

PiperOrigin-RevId: 696286840

Author: IvyZX

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

flaxlib/flaxlib/flaxlib.pyi renamed to flaxlib/flaxlib_cc.pyi

@@ -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