Implement optimized bool_select for primary backends (#3710)

* Implement optimized bool_select for each backend

This commit implements optimized bool_select operations for all Burn backends
to address GitHub issue #3697.

## Changes Made

### NdArray Backend
- Direct boolean array operations using ndarray's select method
- File: crates/burn-ndarray/src/ops/bool_tensor.rs

### Candle Backend
- Leverages Candle's native index_select operation on boolean tensors
- File: crates/burn-candle/src/ops/bool_tensor.rs

### CubeBackend (WGPU/CUDA/ROCm)
- GPU kernel-based implementation using kernel::select
- File: crates/burn-cubecl/src/ops/bool_ops.rs

### Tch Backend (PyTorch)
- Uses PyTorch's efficient index_select_dim operation
- File: crates/burn-tch/src/ops/bool_tensor.rs

The default implementation used inefficient type conversions:
```rust
let int_tensor = B::bool_into_int(tensor);
let selected = B::int_select(int_tensor, dim, indices);
B::int_equal_elem(selected, 1_i32.elem())
```

The optimized implementations eliminate these conversions by using
backend-native boolean selection operations.

Addresses: #3697

* Fix code formatting issues

- Fix function signature formatting in bool_select implementations
- Reorder imports in ndarray backend
- Apply cargo fmt formatting standards

This resolves the CI code-quality check failure.

* Add comprehensive bool select_assign tests

Adds focused tests covering:
- Overlapping indices accumulation behavior
- Complete boolean truth table coverage
- Edge cases (empty indices, multiple true accumulations)
- Verification against original default implementation
- Proof by contradiction for replacement semantics

Tests validate that optimized implementations maintain identical
semantics to the framework's original behavior while providing
performance improvements.

* Add comprehensive test suite for bool_select_assign

- 6 behavior tests covering edge cases and accumulation semantics
- 2 expected failure tests proving replacement semantics would be wrong
- 5 comparison tests validating optimized vs default implementation
- All tests pass confirming OR accumulation behavior is preserved

* Fix formatting

Author: TheDarkchip

crates/burn-candle/src/ops/bool_tensor.rs

@@ -111,6 +111,28 @@ impl<F: FloatCandleElement, I: IntCandleElement> BoolTensorOps<Self> for Candle<
         super::base::flip(tensor, axes)
     }
 
+    fn bool_select(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        indices: IntTensor<Self>,
+    ) -> BoolTensor<Self> {
+        CandleTensor::new(tensor.tensor.index_select(&indices.tensor, dim).unwrap())
+    }
+
+    fn bool_select_assign(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        indices: IntTensor<Self>,
+        value: BoolTensor<Self>,
+    ) -> BoolTensor<Self> {
+        CandleTensor::new(
+            tensor
+                .tensor
+                .index_add(&indices.tensor, &value.tensor, dim)
+                .unwrap(),
+        )
+    }
+
     fn bool_expand(tensor: BoolTensor<Self>, shape: Shape) -> BoolTensor<Self> {
         expand(tensor, shape)
     }

crates/burn-cubecl/src/ops/bool_ops.rs

@@ -106,6 +106,23 @@ where
         expand(tensor, shape)
     }
 
+    fn bool_select(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        indices: IntTensor<Self>,
+    ) -> BoolTensor<Self> {
+        kernel::select::<R, BT, I>(tensor, dim, indices)
+    }
+
+    fn bool_select_assign(
+        tensor: BoolTensor<Self>,
+        dim: usize,
+        indices: IntTensor<Self>,
+        value: BoolTensor<Self>,
+    ) -> BoolTensor<Self> {
+        kernel::select_assign::<R, BT, I>(tensor, dim, indices, value)
+    }
+
     fn bool_flip(tensor: BoolTensor<Self>, axes: &[usize]) -> BoolTensor<Self> {
         kernel::flip::<R, BT, BT>(tensor, axes)
     }

crates/burn-ndarray/src/ops/bool_tensor.rs

@@ -8,7 +8,7 @@ use ndarray::IntoDimension;
 
 // Current crate
 use crate::element::{FloatNdArrayElement, IntNdArrayElement, QuantElement};
-use crate::{NdArray, tensor::NdArrayTensor};
+use crate::{NdArray, execute_with_int_dtype, tensor::NdArrayTensor};
 use crate::{NdArrayDevice, SharedArray};
 
 // Workspace crates
@@ -117,6 +117,40 @@ where
         NdArrayOps::expand(tensor.bool(), shape).into()
     }
 
+    fn bool_select(tensor: NdArrayTensor, dim: usize, indices: NdArrayTensor) -> NdArrayTensor {
+        execute_with_int_dtype!(indices, I, |indices: SharedArray<I>| -> NdArrayTensor {
+            let tensor_bool = tensor.bool();
+            let indices_vec: Vec<usize> = indices
+                .into_iter()
+                .map(|i| i.elem::<i64>() as usize)
+                .collect();
+
+            let selected = tensor_bool.select(ndarray::Axis(dim), &indices_vec);
+            selected.into_shared().into()
+        })
+    }
+
+    fn bool_select_assign(
+        tensor: NdArrayTensor,
+        dim: usize,
+        indices: NdArrayTensor,
+        value: NdArrayTensor,
+    ) -> NdArrayTensor {
+        execute_with_int_dtype!(indices, I, |indices: SharedArray<I>| -> NdArrayTensor {
+            let mut output_array = tensor.bool().into_owned();
+            let value_bool = value.bool();
+
+            for (index_value, index) in indices.into_iter().enumerate() {
+                let index_usize = index.elem::<i64>() as usize;
+                let mut view = output_array.index_axis_mut(ndarray::Axis(dim), index_usize);
+                let value_slice = value_bool.index_axis(ndarray::Axis(dim), index_value);
+                // For boolean tensors, select_assign should use logical OR operation
+                view.zip_mut_with(&value_slice, |a, b| *a = *a || *b);
+            }
+            output_array.into_shared().into()
+        })
+    }
+
     fn bool_flip(tensor: NdArrayTensor, axes: &[usize]) -> NdArrayTensor {
         NdArrayOps::flip(tensor.bool(), axes).into()
     }

crates/burn-tch/src/ops/bool_tensor.rs

@@ -134,6 +134,19 @@ impl<E: TchElement> BoolTensorOps<Self> for LibTorch<E> {
         TchTensor::new(tensor.tensor.argwhere())
     }
 
+    fn bool_select(tensor: TchTensor, dim: usize, indices: TchTensor) -> TchTensor {
+        TchOps::index_select_dim(tensor, dim, indices)
+    }
+
+    fn bool_select_assign(
+        tensor: TchTensor,
+        dim: usize,
+        indices: TchTensor,
+        value: TchTensor,
+    ) -> TchTensor {
+        TchOps::select_assign(tensor, dim, indices, value)
+    }
+
     fn bool_expand(tensor: TchTensor, shape: Shape) -> TchTensor {
         TchOps::expand(tensor, shape)
     }

crates/burn-tensor/src/tests/ops/select.rs

@@ -220,6 +220,244 @@ mod tests {
         output.into_data().assert_eq(&expected, false);
     }
 
+    #[test]
+    fn should_select_assign_bool_overlapping_indices() {
+        // Test accumulation behavior with overlapping indices
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([false, true], &device);
+        let indices = TestTensorInt::from_data([0, 0], &device);
+        let values = TestTensorBool::<1>::from_data([true, false], &device);
+
+        let output = tensor.select_assign(0, indices, values);
+        // Index 0: false OR true OR false = true
+        let expected = TensorData::from([true, true]);
+
+        output.into_data().assert_eq(&expected, false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_false_to_true_case() {
+        // Test false OR true = true
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([false], &device);
+        let indices = TestTensorInt::from_data([0], &device);
+        let values = TestTensorBool::<1>::from_data([true], &device);
+
+        let output = tensor.select_assign(0, indices, values);
+        let expected = TensorData::from([true]);
+
+        output.into_data().assert_eq(&expected, false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_empty_indices() {
+        // Test empty indices array
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true, false, true], &device);
+        let indices = TestTensorInt::<1>::from_data([] as [i32; 0], &device);
+        let values = TestTensorBool::<1>::from_data([] as [bool; 0], &device);
+
+        let output = tensor.select_assign(0, indices, values);
+        let expected = TensorData::from([true, false, true]);
+
+        output.into_data().assert_eq(&expected, false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_true_or_true_accumulation() {
+        // Test multiple true accumulations
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true, false], &device);
+        let indices = TestTensorInt::from_data([0, 0, 0], &device);
+        let values = TestTensorBool::<1>::from_data([true, true, true], &device);
+
+        let output = tensor.select_assign(0, indices, values);
+        let expected = TensorData::from([true, false]);
+
+        output.into_data().assert_eq(&expected, false);
+    }
+
+    #[test]
+    fn should_match_default_implementation_behavior() {
+        // Verify optimized implementation matches original default logic
+        use burn_tensor::backend::Backend;
+
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true, false, true], &device);
+        let indices = TestTensorInt::from_data([0, 1, 0], &device);
+        let values = TestTensorBool::<1>::from_data([false, true, true], &device);
+
+        let optimized_result = tensor
+            .clone()
+            .select_assign(0, indices.clone(), values.clone());
+
+        // Manual default implementation logic
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+
+        optimized_result
+            .into_data()
+            .assert_eq(&default_result.into_data(), false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_overlapping_indices_vs_default() {
+        // Test overlapping indices against default implementation
+        use burn_tensor::backend::Backend;
+
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([false, true], &device);
+        let indices = TestTensorInt::from_data([0, 0], &device);
+        let values = TestTensorBool::<1>::from_data([true, false], &device);
+
+        let optimized_result = tensor
+            .clone()
+            .select_assign(0, indices.clone(), values.clone());
+
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+
+        optimized_result
+            .into_data()
+            .assert_eq(&default_result.into_data(), false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_true_or_true_accumulation_vs_default() {
+        // Test multiple true accumulations against default implementation
+        use burn_tensor::backend::Backend;
+
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true, false], &device);
+        let indices = TestTensorInt::from_data([0, 0, 0], &device);
+        let values = TestTensorBool::<1>::from_data([true, true, true], &device);
+
+        let optimized_result = tensor
+            .clone()
+            .select_assign(0, indices.clone(), values.clone());
+
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+
+        optimized_result
+            .into_data()
+            .assert_eq(&default_result.into_data(), false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_false_to_true_case_vs_default() {
+        // Test false OR true case against default implementation
+        use burn_tensor::backend::Backend;
+
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([false], &device);
+        let indices = TestTensorInt::from_data([0], &device);
+        let values = TestTensorBool::<1>::from_data([true], &device);
+
+        let optimized_result = tensor
+            .clone()
+            .select_assign(0, indices.clone(), values.clone());
+
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+
+        optimized_result
+            .into_data()
+            .assert_eq(&default_result.into_data(), false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_empty_indices_vs_default() {
+        // Test empty indices against default implementation
+        use burn_tensor::backend::Backend;
+
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true, false, true], &device);
+        let indices = TestTensorInt::<1>::from_data([] as [i32; 0], &device);
+        let values = TestTensorBool::<1>::from_data([] as [bool; 0], &device);
+
+        let optimized_result = tensor
+            .clone()
+            .select_assign(0, indices.clone(), values.clone());
+
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+
+        optimized_result
+            .into_data()
+            .assert_eq(&default_result.into_data(), false);
+    }
+
+    #[test]
+    fn should_select_assign_bool_tensor_vs_default() {
+        // Test existing basic case against default implementation
+        use burn_tensor::backend::Backend;
+
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true, false, true], &device);
+        let indices = TestTensorInt::from_data([0, 2], &device);
+        let values = TestTensorBool::<1>::from_data([false, false], &device);
+
+        let optimized_result = tensor
+            .clone()
+            .select_assign(0, indices.clone(), values.clone());
+
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+
+        optimized_result
+            .into_data()
+            .assert_eq(&default_result.into_data(), false);
+    }
+
+    #[test]
+    #[should_panic(expected = "Tensors are not eq")]
+    fn should_fail_if_replacement_semantics_were_used() {
+        // Test that framework uses accumulation, not replacement
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true], &device);
+        let indices = TestTensorInt::from_data([0], &device);
+        let values = TestTensorBool::<1>::from_data([false], &device);
+
+        let output = tensor.select_assign(0, indices, values);
+        let replacement_expected = TensorData::from([false]);
+
+        output.into_data().assert_eq(&replacement_expected, false);
+    }
+
+    #[test]
+    #[should_panic(expected = "Tensors are not eq")]
+    fn should_fail_if_replacement_semantics_were_used_vs_default() {
+        // Test that default implementation also uses accumulation, not replacement
+        use burn_tensor::backend::Backend;
+        let device = Default::default();
+        let tensor = TestTensorBool::<1>::from_data([true], &device);
+        let indices = TestTensorInt::from_data([0], &device);
+        let values = TestTensorBool::<1>::from_data([false], &device);
+
+        let int_tensor = tensor.int();
+        let int_values = values.int();
+        let assigned = int_tensor.select_assign(0, indices, int_values);
+        let default_result = assigned.greater_elem(0);
+        let replacement_expected = TensorData::from([false]);
+
+        default_result
+            .into_data()
+            .assert_eq(&replacement_expected, false);
+    }
+
     #[test]
     fn should_select_with_negative_dim_2d() {
         // Test using negative dimension indexing on 2D tensor