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.

[tiruka]

@laggui

Thank you for your feedback. (sorry for my belated reply)

I understand your point about removing the value check for index_tensor, as invalid values will eventually cause a panic further down the line, likely during the scatter operation.

However, simply removing the check while keeping only the num_classes <= 1 validation (as below codes) still results in a panic during execution when index_tensor contains values greater than or equal to num_classes. This is because the check you suggested removing (index_tensor.clone().greater_equal_elem(num_classes as i32).any().into_scalar()) was the only mechanism to catch this issue early.

    pub(crate) fn one_hot_tensor(
        num_classes: usize,
    ) -> Self {
        let mut check = Self::Ok;
        if num_classes <= 1 {
            check = check.register(
                "One Hot",
                TensorError::new("Can't create a one hot tensor with less then 2 classes"),
            )
        }
        check
    }

Additionally, if this panic occurs outside of the check function, it can lead to a poisoned lock error, which we want to avoid. Given this, do you think we should retain the value check to prevent this scenario, or would you suggest another way to handle this while maintaining performance?

Looking forward to your thoughts!

[laggui]

Just to clarify, the comment was just a thought regarding the value check since it is the only check in place that is not really "static" but rather acts on the values provided to a tensor operation.

Keeping the check makes it more explicit and makes the program panic early, which is good, but requires that all calls to create a one-hot tensor also check the validity of the input as an extra operation, which is not always required if you know the input is valid. So it's a trade-off.

But I do not mind keeping the check just for a better user experience at this time 🙂 it's just something to keep in mind

[tiruka]

@laggui
Thank you! Yes, I also don’t think this is something that needs to be changed immediately, so I believe we are on the same page. I also understand the reasoning behind wanting to improve performance if possible.

However, simply removing the greater check causes test failures, so I believe I need to take a deeper look at the code to properly address this. That was the main purpose of this discussion—I was hoping to get advice on how best to approach this.

First, I wanted to confirm whether the error itself is expected. (From your response, I assume the answer is yes.)
Second, if there is a way to avoid the error, I would like to know. (At this point, it seems that running the check is still necessary.)

Given the above, I think I need to analyze the code more thoroughly. Lastly, if there is a specific part of the code that is the main reason for the performance impact, and if there are any references that might help in avoiding it, could you share them? I would like to use that as a reference to explore possible solutions.

[laggui]

However, simply removing the greater check causes test failures

Yeah that's is specifically because we test for this condition here. If we remove the check, the test can be removed and the panic condition documented. But we can keep it for now as discussed.

Second, if there is a way to avoid the error

Whether we have a check or not, the only way to avoid the error in a user's program is to provide a valid input for one-hot encoding.

Lastly, if there is a specific part of the code that is the main reason for the performance impact, and if there are any references that might help in avoiding it, could you share them?

The simple reason is that the check performs additional tensor operations to validate the values of the provided inputs, whereas a static check (like all the other checks) only perform validation of tensor properties (e.g., shape, rank, etc.). So there is an additional cost to perform this

burn/crates/burn-tensor/src/tensor/api/check.rs

Lines 455 to 459 in 04c5e74

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

(and we also read the bool value of the result)

[tiruka]

Thank you for your reply. As of now we have to keep it as it is and in the future are going to modify it.