PyTorch and Bazel

(Tim Zaman) #1


I’m creating PyTorch C++ extensions and building them with bazel.

The documentation (source) explains how to build your extensions with either setuptools or JIT. However, Bazel likes to take the building into its own hands.

I’d like to get some advise on how to proceed. I currently have two working solutions:

  1. Hacky solution
    Per extension, create a bazel genrule, that just invokes a python build, and set the resulting .so file as the output artifact. This can then be loaded in the code. This leverages all the nice abstractions and build arguments that are set through the torch utilities (torch.utils.cpp_extension.$)

  2. Proper solution
    Create the library through bazel’s cc_library. This is nice, but everything (arguments, flags, includes, directories, etc) need to be set manually.

Is anyone using pytorch extensions with bazel already, or does anyone have any general advice here.