IndicesBoundaryMasker in JAX causes halting problem when used in multi-GPU

The current implementation occasionally causes halting between GPUs when padding is applied to the bmap. This issue primarily arises because the function cannot be JIT-compiled.

In the JAX implementation of IndicesBoundaryMasker, there are several operations, such as conditional statements, that are not supported in JIT-compiled JAX. The previous implementation was JIT-compatible and did not encounter these issues, which should serve as the reference for resolving this problem.