Update nanobind to v2.0

[hawkinsp]

nanobind v2.0 was released today; google/benchmark should upgrade.

A change something like the following should work:

--- a/third_party/benchmark/bindings/python/google_benchmark/benchmark.cc
+++ b/third_party/benchmark/bindings/python/google_benchmark/benchmark.cc
@@ -109,7 +109,7 @@ NB_MODULE(_benchmark, m) {
   using benchmark::Counter;
   nb::class_<Counter> py_counter(m, "Counter");

-  nb::enum_<Counter::Flags>(py_counter, "Flags")
+  nb::enum_<Counter::Flags>(py_counter, "Flags", nb::is_arithmetic())
       .value("kDefaults", Counter::Flags::kDefaults)
       .value("kIsRate", Counter::Flags::kIsRate)
       .value("kAvgThreads", Counter::Flags::kAvgThreads)
@@ -121,7 +121,7 @@ NB_MODULE(_benchmark, m) {
       .value("kAvgIterationsRate", Counter::Flags::kAvgIterationsRate)
       .value("kInvert", Counter::Flags::kInvert)
       .export_values()
-      .def(nb::self | nb::self);
+      .def("__or__", [](Counter::Flags a, Counter::Flags b) { return a | b; });

   nb::enum_<Counter::OneK>(py_counter, "OneK")
       .value("kIs1000", Counter::OneK::kIs1000)

although I don't have time right now to figure out the bazel build updates, etc.

[nicholasjng]

This will require a new release of nanobind-bazel on the BCR, since that's being used to generate the extension.

I haven't gotten around to implement stubgen yet to generate stubs during the Bazel build, but that should not very difficult (albeit I'm not a Bazel expert). I'll find some time to get it done soon.

EDIT: nvm, you already posted an issue. thanks!

[varshneydevansh]

I am trying to do this and did the version change in the MODULE.build file.

 base  devansh   1790_update_nanobind -  bazel build                               [24/05/24| 6:37PM]
Extracting Bazel installation...
Starting local Bazel server and connecting to it...
ERROR: Error computing the main repository mapping: in module dependency chain <root> -> nanobind_bazel@2.0.0: module not found in registries: nanobind_bazel@2.0.0
Computing main repo mapping: 
    Fetching https://bcr.bazel.build/modules/rules_python/0.31.0/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/googletest/1.12.1/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/platforms/0.0.7/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/protobuf/3.19.6/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/apple_support/1.5.0/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/rules_java/7.4.0/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/buildozer/6.4.0.2/MODULE.bazel
    Fetching https://bcr.bazel.build/modules/rules_foreign_cc/0.10.1/MODULE.bazel

Now I understood how it's dependent on the https://github.com/nicholasjng/nanobind-bazel

So, now I have to implement stubgen to generate stubs in the nanobind-bazel as told by @nicholasjng ?

[nicholasjng]

I'm currently working on upgrading to nanobind v2.4.0, but it will take a few more days. Out of curiosity, how do you guys at JAX handle generated stubs for the XLA extensions? Is it fine to just package them directly after running the generating Bazel rule in the wheel build like I do here, or do you prefer checking them into the repository by hand?

[hawkinsp]

I'm currently working on upgrading to nanobind v2.4.0, but it will take a few more days. Out of curiosity, how do you guys at JAX handle generated stubs for the XLA extensions? Is it fine to just package them directly after running the generating Bazel rule in the wheel build like I do here, or do you prefer checking them into the repository by hand?

We hand-author them and check them into the repository.

(At some point, we'll probably explore generated stubs.)