Segmentation fault with simple comparison

[NeilGirdhar]

Description

import jax.numpy as jnp
y = 2.0
jy = jnp.asarray(2.0)
g = {'y': y}
h = {'y': jy}
print(g == h)  # Fault!!

gives

WARNING:2025-06-10 14:14:30,805:jax._src.xla_bridge:791: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
fish: Job 1, 'python a.py' terminated by signal SIGSEGV (Address boundary error)

System info (python version, jaxlib version, accelerator, etc.)

Python 3.13.4
jax==0.6.1
jaxlib==0.6.1

[jakevdp]

Thanks for the report! I'm not able to reproduce this on any of my machines. Can you give us more information about your system architecture and exactly how you installed JAX?

[NeilGirdhar]

❯ lsb_release -a
No LSB modules are available.
Distributor ID:	Ubuntu
Description:	Ubuntu 25.04
Release:	25.04
Codename:	plucky

❯ rm -rf .venv
❯ deactivate
❯ uv venv
Using CPython 3.13.4
Creating virtual environment at: .venv
Activate with: source .venv/bin/activate.fish
❯ uv pip install jax
Resolved 6 packages in 173ms
Installed 6 packages in 17ms
 + jax==0.6.1
 + jaxlib==0.6.1
 + ml-dtypes==0.5.1
 + numpy==2.3.0
 + opt-einsum==3.4.0
 + scipy==1.15.3
❯ source .venv/bin/activate.fish
❯ python a.py
WARNING:2025-06-10 14:31:31,786:jax._src.xla_bridge:791: An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
fish: Job 1, 'python a.py' terminated by signal SIGSEGV (Address boundary error)

[rdyro]

I can repro this with uv's version, but not with pyenv's latest version:

segfaults (version from uv):            3.13.4 (main, Jun  4 2025, 17:37:06) [Clang 20.1.4 ]
does NOT segfault (version from pyenv): 3.13.4 (main, Jun 11 2025, 01:53:24) [GCC 14.2.0]

Stack trace when segfaulting from gdb

Thread 1 "python3" received signal SIGSEGV, Segmentation fault.
0x00007ffff76875d7 in ?? () from /lib/x86_64-linux-gnu/libgcc_s.so.1
(gdb) where
#0  0x00007ffff76875d7 in ?? () from /lib/x86_64-linux-gnu/libgcc_s.so.1
#1  0x00007ffff768949e in _Unwind_Backtrace () from /lib/x86_64-linux-gnu/libgcc_s.so.1
#2  0x00007ffff7dc5cfd in __GI___backtrace (array=<optimized out>, size=128) at ./debug/backtrace.c:78
#3  0x00007ffec549482d in tsl::CurrentStackTrace[abi:cxx11]() () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#4  0x00007ffeca73c34a in xla::RecordGpuCompilerStacktrace() () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#5  0x00007ffec551ed73 in xla::gpu::GpuCompiler::RunBackend(std::unique_ptr<xla::HloModule, std::default_delete<xla::HloModule> >, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&) () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#6  0x00007ffec54e86f2 in xla::Service::BuildExecutable(xla::HloModuleProto const&, std::unique_ptr<xla::HloModuleConfig, std::default_delete<xla::HloModuleConfig> >, xla::Backend*, stream_executor::StreamExecutor*, xla::Compiler::CompileOptions const&, bool) () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#7  0x00007ffec54b0357 in xla::LocalService::CompileExecutables(xla::XlaComputation const&, absl::lts_20230802::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#8  0x00007ffec54ad334 in xla::LocalClient::Compile(xla::XlaComputation const&, absl::lts_20230802::Span<xla::Shape const* const>, xla::ExecutableBuildOptions const&) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#9  0x00007ffec5412319 in xla::PjRtStreamExecutorClient::CompileInternal(xla::XlaComputation const&, std::vector<xla::Shape const*, std::allocator<xla::Shape const*> > const&, std::function<absl::lts_20230802::StatusOr<std::pair<std::vector<xla::Shape, std::allocator<xla::Shape> >, xla::Shape> > (xla::HloModule const&)>, xla::CompileOptions, bool) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#10 0x00007ffec5413b67 in xla::PjRtStreamExecutorClient::Compile(mlir::ModuleOp, xla::CompileOptions, bool) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#11 0x00007ffec5414856 in xla::PjRtStreamExecutorClient::CompileAndLoad(mlir::ModuleOp, xla::CompileOptions) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#12 0x00007ffec53d18ef in xla::StreamExecutorGpuClient::CompileAndLoad(mlir::ModuleOp, xla::CompileOptions) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#13 0x00007ffec538276a in std::__detail::__variant::__gen_vtable_impl<true, std::__detail::__variant::_Multi_array<absl::lts_20230802::StatusOr<std::unique_ptr<xla::PjRtLoadedExecutable, std::default_delete<xla::PjRtLoadedExecutable> > > (*)(pjrt::PJRT_Client_Compile(PJRT_Client_Compile_Args*)::$_0&&, std::variant<mlir::OwningOpRef<mlir::ModuleOp>, xla::XlaComputation>&)>, std::tuple<std::variant<mlir::OwningOpRef<mlir::ModuleOp>, xla::XlaComputation>&>, std::integer_sequence<unsigned long, 0ul> >::__visit_invoke(pjrt::PJRT_Client_Compile(PJRT_Client_Compile_Args*)::$_0&&, std::variant<mlir::OwningOpRef<mlir::ModuleOp>, xla::XlaComputation>&) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#14 0x00007ffec5371b54 in pjrt::PJRT_Client_Compile(PJRT_Client_Compile_Args*) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jax_plugins/xla_cuda12/xla_cuda_plugin.so
#15 0x00007fff48a6169d in xla::InitializeArgsAndCompile(xla::PjRtCApiClient*, PJRT_Api const*, PJRT_Client*, xla::CompileOptions const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#16 0x00007fff48a61e06 in xla::PjRtCApiClient::CompileAndLoad(mlir::ModuleOp, xla::CompileOptions) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#17 0x00007fff508599f8 in xla::ifrt::PjRtLoadedExecutable::Create(xla::ifrt::PjRtCompatibleClient*, mlir::ModuleOp, xla::CompileOptions, std::vector<tsl::RCReference<xla::ifrt::LoadedHostCallback>, std::allocator<tsl::RCReference<xla::ifrt::LoadedHostCallback> > >, xla::ifrt::RCReferenceWrapper<xla::ifrt::DeviceList>) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#18 0x00007fff508546f7 in xla::ifrt::PjRtCompiler::CompileAndLoad(std::unique_ptr<xla::ifrt::Program, std::default_delete<xla::ifrt::Program> >, std::unique_ptr<xla::ifrt::CompileOptions, std::default_delete<xla::ifrt::CompileOptions> >) () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#19 0x00007fff4ef49f80 in xla::PyClient::CompileAndLoadIfrtProgram(xla::nb_class_ptr<xla::PyClient>, std::unique_ptr<xla::ifrt::Program, std::default_delete<xla::ifrt::Program> >, std::unique_ptr<xla::ifrt::CompileOptions, std::default_delete<xla::ifrt::CompileOptions> >) () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#20 0x00007fff4ef4afab in xla::PyClient::CompileAndLoad(xla::nb_class_ptr<xla::PyClient>, std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> >, xla::ifrt::RCReferenceWrapper<xla::ifrt::DeviceList>, xla::CompileOptions, std::vector<nanobind::capsule, std::allocator<nanobind::capsule> >) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#21 0x00007fff4ef55248 in nanobind::detail::func_create<false, true, xla::PyClient::RegisterPythonTypes(nanobind::module_&)::$_2, xla::nb_class_ptr<xla::PyLoadedExecutable>, xla::nb_class_ptr<xla::PyClient>, nanobind::bytes, jax::PyDeviceList&, xla::CompileOptions, std::vector<nanobind::capsule, std::allocator<nanobind::capsule> >, 0ul, 1ul, 2ul, 3ul, 4ul, nanobind::scope, nanobind::name, nanobind::is_method, nanobind::arg, nanobind::arg, nanobind::arg_v, nanobind::arg_v>(xla::PyClient::RegisterPythonTypes(nanobind::module_&)::$_2&&, xla::nb_class_ptr<xla::PyLoadedExecutable> (*)(xla::nb_class_ptr<xla::PyClient>, nanobind::bytes, jax::PyDeviceList&, xla::CompileOptions, std::vector<nanobind::capsule, std::allocator<nanobind::capsule> >), std::integer_sequence<unsigned long, 0ul, 1ul, 2ul, 3ul, 4ul>, nanobind::scope const&, nanobind::name const&, nanobind::is_method const&, nanobind::arg const&, nanobind::arg const&, nanobind::arg_v const&, nanobind::arg_v const&)::{lambda(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*)#1}::__invoke(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*) () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#22 0x00007fff50812b31 in nanobind::detail::nb_func_vectorcall_complex(_object*, _object* const*, unsigned long, _object*) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#23 0x00007fff54d8640c in nanobind::detail::nb_bound_method_vectorcall(_object*, _object* const*, unsigned long, _object*) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/cpu/_lapack.so
#24 0x0000000001a2467f in _PyEval_EvalFrameDefault ()
#25 0x0000000001ab6750 in PyObject_Vectorcall ()
#26 0x00007fff50809d3d in nanobind::detail::obj_vectorcall(_object*, _object* const*, unsigned long, _object*, bool) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#27 0x00007fff51f2d903 in nanobind::object nanobind::detail::api<nanobind::handle>::operator()<(nanobind::rv_policy)1, nanobind::object&, nanobind::detail::args_proxy, nanobind::detail::kwargs_proxy>(nanobind::object&, nanobind::detail::args_proxy&&, nanobind::detail::kwargs_proxy&&) const ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#28 0x00007fff51f2c8a6 in jax::WeakrefLRUCache::Call(nanobind::object, nanobind::args, nanobind::kwargs) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#29 0x00007fff51f2f583 in nanobind::detail::func_create<false, true, nanobind::cpp_function_def<jax::WeakrefLRUCache, nanobind::object, jax::WeakrefLRUCache, nanobind::object, nanobind::args, nanobind::kwargs, nanobind::scope, nanobind::name, nanobind::is_method, nanobind::lock_self>(nanobind::object (jax::WeakrefLRUCache::*)(nanobind::object, nanobind::args, nanobind::kwargs), nanobind::scope const&, nanobind::name const&, nanobind::is_method const&, nanobind::lock_self const&)::{lambda(jax::WeakrefLRUCache*, nanobind::object, nanobind::args, nanobind::kwargs)#1}, nanobind::object, jax::WeakrefLRUCache*, nanobind::object, nanobind::args, nanobind::kwargs, 0ul, 1ul, 2ul, 3ul, nanobind::scope, nanobind::name, nanobind::is_method, nanobind::lock_self>(nanobind::cpp_function_def<jax::WeakrefLRUCache, nanobind::object, jax::WeakrefLRUCache, nanobind::object, nanobind::args, nanobind::kwargs, nanobind::scope, nanobind::name, nanobind::is_method, nanobind::lock_self>(nanobind::object (jax::WeakrefLRUCache::*)(nanobind::object, nanobind::args, nanobind::kwargs), nanobind::scope const&, nanobind::name const&, nanobind::is_method const&, nanobind::lock_self const&)::{lambda(jax::WeakrefLRUCache*, nanobind::object, nanobind::args, nanobind::kwargs)#1}&&, nanobind::object (*)(jax::WeakrefLRUCache*, nanobind::object, nanobind::args, nanobind::kwargs), std::integer_sequence<unsigned long, 0ul, 1ul, 2ul, 3ul>, nanobind::scope const&, nanobind::name const&, nanobind::is_method const&, nanobind::lock_self const&)::{lambda(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*)#1}::__invoke(void*, _object**, unsigned char*, nanobind::rv_policy, nanobind::detail::cleanup_list*) () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#30 0x00007fff50812b31 in nanobind::detail::nb_func_vectorcall_complex(_object*, _object* const*, unsigned long, _object*) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#31 0x0000000001ae92ae in _PyObject_Call_Prepend ()
#32 0x0000000001ae91f3 in slot_tp_call ()
#33 0x0000000001a19add in _PyEval_EvalFrameDefault ()
#34 0x0000000001ab6750 in PyObject_Vectorcall ()
#35 0x00007fff489e35b6 in jax::(anonymous namespace)::PjitFunction::Call(nanobind::handle, _object* const*, unsigned long, _object*) ()
   from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#36 0x00007fff489e161f in PjitFunction_tp_vectorcall () from /usr/local/google/home/rdyro/test/lib/python3.13/site-packages/jaxlib/libjax_common.so
#37 0x0000000001a27f2d in _PyEval_EvalFrameDefault ()
#38 0x0000000001a0596d in slot_tp_richcompare ()
#39 0x0000000001cb6a07 in PyObject_RichCompareBool.warm ()
#40 0x0000000000000003 in ?? ()
#41 0x0000000001948218 in _PyRuntime ()
#42 0x0000000000000008 in ?? ()
#43 0x0000000000000000 in ?? ()
(gdb)

CUDA info:

Wed Jun 11 01:47:58 2025
+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.216.03             Driver Version: 535.216.03   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|

[NeilGirdhar]

Interesting! Any idea why that might be?

[hawkinsp]

I think calling tsl::CurrentStackTrace is not a safe thing to do except as a debugging tool.

Here it's being used as part of some metrics collection code in XLA... but I don't think anyone can even consume those metrics in our public JAX build. So it's doing nothing useful and we should just disable that code.

[hawkinsp]

Already fixed in the latest release (0.6.2)!