Editable jaxlib build

[yhtang]

This PR adds an --editable option to the build/build.py script. When set, the build script will

1. Upon compilation completion, copy the jaxlib source folder and the compiled .so files from the Bazel sandbox to the output path (dist by default)
2. print a message prompting users to use pip install -e to install jaxlib in-place at the output path.

This way, we only need to run pip install -e once, and subsequent Bazel builds with --editable will automatically update the installed jaxlib package (e.g., in a virtual environment).

Things worth consideration:

• Would it be better to copy the compiled .so files directly back into the jaxlib source folder and run pip install -e there? This way, rebuilding is not needed if changes only apply to Python files.

This PR is adapted from the Alpa project as part of a rebasing effort, with contributions from @merrymercy, @zhisbug, @ZYHowell, @zhuohan123 among many others.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[yhtang]

@mjsML @nouiz for viz.

[yhtang]

Updated commit message's co-author list after communicating with the original author @merrymercy

[hawkinsp]

Things worth consideration:

• Would it be better to copy the compiled .so files directly back into the jaxlib source folder and run pip install -e there? This way, rebuilding is not needed if changes only apply to Python files.

Perhaps, but I think not. In general it's nice to keep built artifacts out of the source tree.