Problems since using bc algorithms

[Excurrybang]

Hello brax team,

thanks for your implementation of bc, there are some code problems I found when using the bc. I tried to fix them with AI, but I think you may have better solution for these problems. most of them are facing jax compile errors.

thanks, Tau