The role of check! in numeric.rs

[tiruka]

Purpose

Relating to #2613 (comment), so let me mention @laggui

I would like to check the role of check! macro and related check functions in crates/burn-tensor/src/tensor/api/check.rs
I suppose it is necessary to prevent poisoned lock issues instead of native errors.
If my guess is right, should I keep the one hot tensor check function?

Background

I tried to remove the parts of one_hot_tensor as below and run the command cargo xtask validate or cargo test -p burn-router --lib.

        if index_tensor
            .clone()
            .greater_equal_elem(num_classes as i32)
            .any()
            .into_scalar()
        {

Then, I come across poisoned lock error as below.

thread 'tests::tensor::one_hot::tests::float_one_hot_should_panic_when_index_exceeds_number_of_classes' panicked at /Users/tiruka/Desktop/oss-fork/burn/crates/burn-ndarray/src/ops/base.rs:367:23:
ndarray: index [0, 5] is out of bounds for array of shape [1, 5]
thread 'tests::tensor::one_hot::tests::float_one_hot_should_panic_when_index_exceeds_number_of_classes' panicked at crates/burn-router/src/runner.rs:1233:29:
called `Result::unwrap()` on an `Err` value: "poisoned lock: another task failed inside"
stack backtrace:
   0:        0x1064e3764 - <std::sys::backtrace::BacktraceLock::print::DisplayBacktrace as core::fmt::Display>::fmt::h7adefaad4a31afc0
   1:        0x10650007c - core::fmt::write::h381c0b0ce6ab972a
   2:        0x1064e0f24 - std::io::Write::write_fmt::h9610cdd27cd3672c
   3:        0x1064e3618 - std::sys::backtrace::BacktraceLock::print::h8baf33e22611de71
   4:        0x1064e4984 - std::panicking::default_hook::{{closure}}::h369c7295ef58c5b1
   5:        0x1064e4794 - std::panicking::default_hook::h50746358288a9d6a
   6:        0x105d235c8 - test::test_main::{{closure}}::h167d719e9b545fe1
   7:        0x1064e5174 - std::panicking::rust_panic_with_hook::h7d795911432661cb
   8:        0x1064e4dec - std::panicking::begin_panic_handler::{{closure}}::h36f15310ecbde379
   9:        0x1064e3c28 - std::sys::backtrace::__rust_end_short_backtrace::heed121414170e0c7
  10:        0x1064e4aa4 - _rust_begin_unwind
  11:        0x106559590 - core::panicking::panic_fmt::h17b1b80ec02ffd19
  12:        0x106559980 - core::result::unwrap_failed::h6de8f20130c8e837
  13:        0x104d9e838 - <burn_router::runner::Runner<B> as burn_router::client::base::RunnerClient>::register_orphan::h400c7060a8faaff4
  14:        0x105612614 - <burn_router::types::duo::MultiRunnerClient<B1,B2> as burn_router::client::base::RunnerClient>::register_orphan::h58ca8091cb96ce02
  15:        0x10571fbfc - <burn_router::tensor::RouterTensor<C> as core::ops::drop::Drop>::drop::h5e34a7dc836967a0
  16:        0x1056cccf8 - core::ptr::drop_in_place<burn_router::tensor::RouterTensor<burn_router::types::duo::MultiRunnerClient<burn_ndarray::backend::NdArray<f32,i32>,burn_jit::backend::JitBackend<cubecl_wgpu::runtime::WgpuRuntime,f32,i32,u32>>>>::h8d9b05171dba0898
  17:        0x104ca8030 - burn_router::ops::op_float::<impl burn_tensor::tensor::ops::tensor::FloatTensorOps<burn_router::backend::BackendRouter<R>> for burn_router::backend::BackendRouter<R>>::float_scatter::h27d77222a5ca62c8
  18:        0x104df6040 - <burn_tensor::tensor::api::kind::Float as burn_tensor::tensor::api::numeric::Numeric<B>>::scatter::ha93ec77e74976daa
  19:        0x105537a10 - burn_tensor::tensor::api::numeric::<impl burn_tensor::tensor::api::base::Tensor<B,_,K>>::scatter::hcc2ddc1bf05cb4c9
  20:        0x105525c40 - burn_tensor::tensor::api::numeric::<impl burn_tensor::tensor::api::base::Tensor<B,_,K>>::one_hot_fill::h58cef82a769bd5d3
  21:        0x105536cb4 - burn_tensor::tensor::api::numeric::<impl burn_tensor::tensor::api::base::Tensor<B,_,K>>::one_hot::h47526580e524cba8
  22:        0x104d43384 - burn_router::tests::tensor::one_hot::tests::float_one_hot_should_panic_when_index_exceeds_number_of_classes::h6dc463e59e8ade2a
  23:        0x1051a956c - burn_router::tests::tensor::one_hot::tests::float_one_hot_should_panic_when_index_exceeds_number_of_classes::{{closure}}::hca866cccce5f7429
  24:        0x10567940c - core::ops::function::FnOnce::call_once::h686c71ceeaf6348a
  25:        0x105d2709c - test::__rust_begin_short_backtrace::he48c7882cfe7d316
  26:        0x105d26a4c - test::run_test::{{closure}}::h7a78ef96c088d027
  27:        0x105cf6d00 - std::sys::backtrace::__rust_begin_short_backtrace::hf59a09bc33aca1d1
  28:        0x105cf9ca0 - core::ops::function::FnOnce::call_once{{vtable.shim}}::h23301f968b1dfa27
  29:        0x1064e864c - std::sys::pal::unix::thread::Thread::new::thread_start::h44ebc83f14222f58
  30:        0x18bb42f94 - __pthread_joiner_wake
thread 'tests::tensor::one_hot::tests::float_one_hot_should_panic_when_index_exceeds_number_of_classes' panicked at core/src/panicking.rs:229:5:

The same errors can be triggered when I remove other check functions like pub fn scatter(self, dim: usize, indices: Tensor<B, D, Int>, values: Self) -> Self {...}. Therefore, I do not think it is only limited to one_hot function.
My guess is these kinds of checks are necessary here, so I should keep checks in one hot as it is.

Reproduction

Use and clone this repository and run cargo test -p burn-router --lib.
main...tiruka:burn:feature/delete-old-one-hot-checker

I may not include some necessary information, in that case, please let me know.

[laggui]

Well first of all, the behavior is triggered in the test because we specifically test for that panic condition. So if you remove the check but leave the test, it will fail. And because we don't have an explicit check anymore for the content of the source tensor you want to convert to its one-hot representation, the failure happens further down the line.

The checks are valid, and there for a reason. But pretty much none of them operate on the content/values of the tensor, instead on some other "static" conditions like the shape, rank, etc.

My comment was specifically directed at the value check. Performing that check in the op has some performance impact, so there is a valid point to be made about its removal. As you can see, the program will panic anyway if you have invalid values, it's just that the panic message is not as clear 😅 And regarding the changes linked, I would not remove the entire check. We can keep the check on the number of classes, but remove the value check.