Proper bounds checking for CubeCL kernels in Burn crates

[AsherJingkongChen]

Describe the bug

Just as the title says. For details, please follow the comments under:

• Refactor for wgpu v23 compatibility #2435

To Reproduce

1. Clone or checkout the PR branch:

   a. Clone

   git clone --branch refactor/wgpu-v23 --depth 1 https://github.com/AsherJingkongChen/burn.git

   b. Checkout

   git remote add pr-2435 https://github.com/AsherJingkongChen/burn.git
   git fetch pr-2435 refactor/wgpu-v23
   git checkout -b refactor/wgpu-v23 pr-2435/refactor/wgpu-v23

2. Run some tests on macOS with Metal:

   a. Test the specific package

   cargo test -p burn-jit

   b. Test for CI

   cargo xtask test --ci

Expected behavior

• All errors during tests on macOS are resolved
• Bounds checking works as expected, it is recommended to add some new tests.

Screenshots

You can check CI statuses on macOS in #2435.

Additional context

This issue blocks #2435.