Modification to Layout Assignment Pass

[danteisalive]

I want to make couple of change to the layout assignment pass in XLA and measure its performance impact on the run time. I build and install Jaxlib with a local XLA repo with these below command lines and I can confirm that JAX is working perfectly fine on Google TPU VMs.

python build/build.py --bazel_options=--override_repository=xla=~/xla
pip install ~/jax/dist/jaxlib-0.4.24.dev20231229-cp311-cp311-manylinux2014_x86_64.whl --force-reinstall

Now, even when I add a simple LOG(INFO) to the layout assignment pass RUN function, it doesn't print anything although I can see that Bazel finds my changes to the local XLA repo, and it compiles them. I can also see other INFOs printed to the terminal. What I'm doing wrong here?