Have anyone faced this issue in print

[jimmylihui]

F1207 17:26:54.735655  190300 gemm_rewriter_triton.cc:715] Check failed: dst_remaining_size == 1 (2 vs. 1) 
*** Check failure stack trace: ***
    @     0x7fa0f3b38e6e  absl::lts_20230802::log_internal::LogMessage::Flush()
    @     0x7fa0f3b38f79  absl::lts_20230802::log_internal::LogMessageFatal::~LogMessageFatal()
    @     0x7fa0eb64dcfe  xla::gpu::(anonymous namespace)::FusionContext::HandleBitcast()
    @     0x7fa0ee959d88  xla::gpu::(anonymous namespace)::FusionContext::HandleInstruction()
    @     0x7fa0ee95cb44  xla::gpu::(anonymous namespace)::FusionContext::AnalyzeForFusionImpl()
    @     0x7fa0ee95d453  xla::gpu::(anonymous namespace)::FusionContext::AnalyzeForFusion()
    @     0x7fa0ee95f59a  xla::gpu::(anonymous namespace)::FusionContext::TryToFuseWithInputsRecursively()
    @     0x7fa0ee960219  xla::gpu::(anonymous namespace)::FuseDot()
    @     0x7fa0ee961bba  xla::gpu::(anonymous namespace)::GemmRewriterTritonVisitor::HandleDot()
    @     0x7fa0f389b600  xla::HloInstruction::Visit<>()
    @     0x7fa0f38c11cf  xla::PostOrderDFS<>()
    @     0x7fa0f38c1fef  xla::HloInstruction::Accept<>()
    @     0x7fa0eb9b16a9  xla::HloComputation::Accept<>()
    @     0x7fa0ee951d8d  xla::gpu::GemmRewriterTriton::Run()
    @     0x7fa0f16dcb45  xla::HloPassPipeline::RunPassesInternal<>()
    @     0x7fa0f16dd937  xla::HloPassPipeline::Run()
    @     0x7fa0ecb319e8  xla::gpu::GpuCompiler::OptimizeHloPostLayoutAssignment()
    @     0x7fa0eca828c4  xla::gpu::NVPTXCompiler::OptimizeHloPostLayoutAssignment()
    @     0x7fa0ecb3cb1b  xla::gpu::GpuCompiler::OptimizeHloModule()
    @     0x7fa0ecb41431  xla::gpu::GpuCompiler::RunHloPasses()
    @     0x7fa0eca71129  xla::Service::BuildExecutable()
    @     0x7fa0ec82affd  xla::LocalService::CompileExecutables()
    @     0x7fa0ec825cd2  xla::LocalClient::Compile()
    @     0x7fa0ec7e4acc  xla::PjRtStreamExecutorClient::Compile()
    @     0x7fa0ec7c09ef  xla::StreamExecutorGpuClient::Compile()
    @     0x7fa0ec7f7d5a  xla::PjRtStreamExecutorClient::Compile()
    @     0x7fa0ec7123bf  xla::ifrt::PjRtLoadedExecutable::Create()
    @     0x7fa0ec708a64  xla::ifrt::PjRtCompiler::Compile()
    @     0x7fa0ebc8d762  xla::PyClient::Compile()
    @     0x7fa0eb9be8b3  pybind11::detail::argument_loader<>::call_impl<>()
    @     0x7fa0eb9bed60  pybind11::cpp_function::initialize<>()::{lambda()#3}::_FUN()
    @     0x7fa0eb973e98  pybind11::cpp_function::dispatcher()
    @           0x4fcb97  cfunction_call

This error appear when I try to print:

jjax.debug.print("0")

if I delete print everything is fine

[mattjj]

I haven't seen that one before... what version of jax and (more importantly) jaxlib? (Trying an upgrade is always a good first step.)

Any chance at sharing a repro?

[jimmylihui]

Hi,
Here is repository:https://github.com/jimmylihui/deer. just run demo2 in deer folder.
Since I update the version. There is another issue happens:

2023-12-07 22:30:27.651374: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 4: 1.01222, expected 0.662715
2023-12-07 22:30:27.651457: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 5: 0.686769, expected 0.480389
2023-12-07 22:30:27.651477: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 6: 0.981596, expected 0.480855
2023-12-07 22:30:27.651500: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 10: 0.480359, expected 0.79812
2023-12-07 22:30:27.651516: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 13: 0.211555, expected 0.851442
2023-12-07 22:30:27.651532: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 14: 0.795369, expected 0.171261
2023-12-07 22:30:27.651547: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 16: 0.797591, expected 0.328714
2023-12-07 22:30:27.651562: E external/xla/xla/service/gpu/buffer_comparator.cc:143] Difference at 17: 0.466938, expected 0.138692
E1207 22:30:27.651576   12035 triton_autotuner.cc:766] Results do not match the reference. This is likely a bug/unexpected loss of precision.
jax.errors.SimplifiedTraceback: For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/ssddata/data/jiahuili/miniconda3/envs/myenv/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/ssddata/data/jiahuili/miniconda3/envs/myenv/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/ssddata/data/jiahuili/.vscode-server/extensions/ms-python.python-2022.20.2/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/__main__.py", line 39, in <module>
    cli.main()
  File "/ssddata/data/jiahuili/.vscode-server/extensions/ms-python.python-2022.20.2/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 430, in main
    run()
  File "/ssddata/data/jiahuili/.vscode-server/extensions/ms-python.python-2022.20.2/pythonFiles/lib/python/debugpy/adapter/../../debugpy/launcher/../../debugpy/../debugpy/server/cli.py", line 284, in run_file
    runpy.run_path(target, run_name="__main__")
  File "/ssddata/data/jiahuili/.vscode-server/extensions/ms-python.python-2022.20.2/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 321, in run_path
    return _run_module_code(code, init_globals, run_name,
  File "/ssddata/data/jiahuili/.vscode-server/extensions/ms-python.python-2022.20.2/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 135, in _run_module_code
    _run_code(code, mod_globals, init_globals,
  File "/ssddata/data/jiahuili/.vscode-server/extensions/ms-python.python-2022.20.2/pythonFiles/lib/python/debugpy/_vendored/pydevd/_pydevd_bundle/pydevd_runpy.py", line 124, in _run_code
    exec(code, run_globals)
  File "/ssddata/data/jiahuili/deer/deer/demo_2.py", line 158, in <module>
    main()
  File "/ssddata/data/jiahuili/deer/deer/demo_2.py", line 126, in main
    outputs = eval_gru(carry,x,y, parameter,Linear,structure, method="deer")
jaxlib.xla_extension.XlaRuntimeError: INTERNAL: All algorithms tried for %dot.65 = f32[5,2,2]{2,1,0} dot(f32[5,2,2]{2,1,0} %slice.47, f32[5,2,2]{2,1,0} %slice.48), lhs_batch_dims={0}, lhs_contracting_dims={2}, rhs_batch_dims={0}, rhs_contracting_dims={1}, metadata={op_name="jit(eval_gru)/jit(main)/while/body/dot_general[dimension_numbers=(((2,), (1,)), ((0,), (0,))) precision=None preferred_element_type=float32]" source_file="/ssddata/data/jiahuili/deer/deer/seq1d_2.py" source_line=197} failed. Falling back to default algorithm.  Per-algorithm errors:
  Results do not match the reference. This is likely a bug/unexpected loss of precision.

[nimashoghi]

I'm having a problem very similar to @jimmylihui's problem above.

I tried installing older versions, and all JAX versions before 0.4.21 work perfectly, but >= 0.4.21 ends up giving me that error (in my case, the operation it complains about is a large einsum which effectively does a matrix multiplication of some input and a precomputed matrix).

[andylolu2]

Had the same problem (Results do not match the reference. This is likely a bug/unexpected loss of precision) on an einsum. Upgrading from 0.4.23 to 0.4.24 fixed it.