Failing Unit Tests

[Heidar-An]

Machine = WSL, Ryzen 7 5800H, RTX 3070 Mobile.
When I run pytest src/metrax a selection of unit tests fail for me. I've attached the error output here:

errors.txt

More than likely the issue is something to do with my machine but wanted to send here just in case.

[rni418]

@jshin1394 - Do you think it makes sense for us to use golden values generated from the TF methods instead of relying on the output of the function call?

[jshin1394]

Seems like this is the error:

2025-05-30 18:16:48.202817: W tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.cc:40] 'cuModuleLoadData(&module, data)' failed with 'CUDA_ERROR_UNSUPPORTED_PTX_VERSION'

2025-05-30 18:16:48.202888: W tensorflow/compiler/mlir/tools/kernel_gen/tf_gpu_runtime_wrappers.cc:40] 'cuModuleGetFunction(&function, module, kernel_name)' failed with 'CUDA_ERROR_INVALID_HANDLE'

2025-05-30 18:16:48.203674: W tensorflow/core/framework/op_kernel.cc:1844] INTERNAL: 'cuLaunchKernel(function, gridX, gridY, gridZ, blockX, blockY, blockZ, 0, reinterpret_cast<CUstream>(stream), params, nullptr)' failed with 'CUDA_ERROR_INVALID_HANDLE'

The error suggests that there is a mismatch in compatibility between your installed TensorFlow version and your NVIDIA driver/CUDA setup and therefore requires a driver update. @Heidar-An were you able to call tensorflow correctly before? Did this error start happening only after you tried running Metrax code?

We used to have hard-coded values generated from the TF methods for tests. But we migrated to calling actual function call in order to be more accurate and robust. Basically, under the same environment, JAX metrics should match keras or tensorflow metric values.