Move generic tests from burn-flex into burn-backend-tests

[antimora]

Proposal

Move generic (not flex-specific) tests out of crates/burn-flex/src/ops/*.rs and into crates/burn-backend-tests/tests/tensor/..., so the shared suite runs them against every backend instead of only burn-flex.

Motivation

crates/burn-flex/src/ops/ currently has ~500 #[test] functions spread across the ops modules. A sizable chunk are plain "does relu work on a 1d f32 tensor" style smoke tests that test backend-agnostic behavior. They live next to the implementation for historical reasons (flex is new and it was convenient during development), but they do not exercise anything specific to FlexTensor's layout system, and the same tests in burn-backend-tests would catch the same regressions while also covering ndarray, wgpu, cuda, etc.

Three costs of the current state:

1. Coverage gap. A bug in the shared dispatch logic (e.g. burn-flex: bool binary ops (and/or) don't broadcast #4771, burn-flex: attention_naive rejects ONNX-compliant broadcasted mask/bias #4772) can sail past backend-tests because it happens to only land in burn-flex unit tests, and backend-tests has no equivalent coverage.
2. Duplication risk. When a new bug is filed, it's natural to add the regression test at the fix site (in burn-flex). That leaves the other backends uncovered. I hit exactly this pattern when fixing burn-flex: bool binary ops (and/or) don't broadcast #4771 and ended up writing the tests twice before dedup.
3. Test-run velocity. cargo test -p burn-flex is bloated with tests that don't need a flex-specific build to run.

What counts as "flex-specific"

Should stay in burn-flex:

• Tests that exercise FlexTensor's stride system directly: test_binary_mul_non_contiguous_transposed, test_binary_add_non_contiguous_narrowed, test_bool_and_flipped, test_bool_and_both_flipped, test_bool_into_int_flipped_2d, etc. These probe negative/permuted/narrowed strides and how the op handles them - something the backend-tests TestTensor API cannot express directly.
• Tests that pin SIMD/Rayon chunk boundaries: test_softmax_simd_body_row, test_softmax_multi_chunk_rayon. These guard against off-by-one bugs at the boundaries of SIMD lanes or worker chunks, which only exist in flex.
• Tests that pin flex-specific dtype storage: test_bool_into_int_u8 (checks BoolStore::U8 backing).
• Tests for flex-internal helpers (layout, strided iterator, SIMD intrinsics).

Should move to burn-backend-tests:

• Plain arithmetic smoke tests: test_binary_add_contiguous_f32/f64, test_binary_single_element, test_binary_1d_tensor, test_binary_3d_tensor, test_scalar_add_f32, etc.
• Plain activation tests: test_relu, test_sigmoid, test_gelu, test_leaky_relu, test_softmax_1d, test_softmax_2d_last_axis, test_layer_norm_* (the ones without SIMD-boundary intent).
• Plain float/int smoke tests in float.rs: test_add_tensors, test_mul_tensors, test_exp_tensor_api, test_log_tensor_api, etc.
• Any #[test] that could be rewritten with TestTensor/TestTensorInt/TestTensorBool without loss of intent.

Rule of thumb: if you can write the test against burn-tensor's public API without touching FlexTensor internals, it belongs in backend-tests.

Scope

This is not a one-shot migration. It can be done file by file as a series of small PRs. A good starting target is the ops with the largest test counts that look mostly-generic, e.g. float.rs (29 tests), activation.rs (26), binary.rs (38). Each PR should:

• Move the generic tests into the matching burn-backend-tests/tests/tensor/.../ops/*.rs file (create if missing).
• Rewrite from FlexTensor::from_data(...) / Flex::bool_and(...) style to TestTensorBool::<D>::from_data(...) / .bool_and(...) style.
• Leave flex-specific tests in place, and add a short comment block at the top of the remaining #[cfg(test)] mod tests explaining what stays and why, so future contributors don't re-add generic tests here.
• Verify the tests pass against at least two backends via feature flags (e.g. cargo test -p burn-backend-tests --no-default-features --features flex,std and ... --features ndarray,std).

Related

• Migrate remaining TestBackend ndarray references to burn-flex #4762 (migrating TestBackend references to burn-flex) is adjacent but different: that's about which backend backend-tests runs against by default, this is about which tests live where.
• Surfaced while working on burn-flex: bool binary ops (and/or) don't broadcast #4771.