[XLA] compilation fails with InvalidArgumentError: Dynamic inferencing on custom call goo.gle/debugstr

[jasminetrail]

tf.function(jit_compile=True) fails to compile a model where the output of tf.linalg.qr is passed directly to tf.sequence_mask. The same program executes successfully in eager mode, but XLA compilation fails with an InvalidArgumentError.

During compilation, XLA reports:

Dynamic inferencing on custom call ProductOfElementaryHouseholderReflectors is not supported while evaluating input 1 of Range operator as a compile-time constant.

The failure occurs because tf.sequence_mask requires determining the maximum sequence length (used by an internal Range operation), and XLA attempts to perform compile-time inference from the result of the QR decomposition. The QR operation is lowered to the custom call ProductOfElementaryHouseholderReflectors, which does not support the required dynamic inference, causing compilation to fail. Here is the gist

Minimal Code

import traceback
import tensorflow as tf

tf.random.set_seed(5415)
t48875 = tf.random.normal([3, 32, 32], dtype=tf.float16)

def model_5415(t48875):
    t48876 = tf.linalg.qr(t48875)[0]
    t48877 = tf.sequence_mask(t48876, dtype=tf.int64)
    return t48877

_args = (t48875,)

print("=== EAGER ===")
try:
    eager_out = model_5415(*_args)
    print("Eager Output Succeeded")
except Exception:
    print(traceback.format_exc())

print("\n=== XLA ===")
try:
    xla_out = tf.function(model_5415, jit_compile=True)(*_args)
    print("Compiler Output Succeeded")
except Exception:
    print(traceback.format_exc())

Error logs

=== EAGER ===
Eager Output Succeeded

=== XLA ===
Traceback (most recent call last):
  File "/tmp/ipykernel_2124/288000946.py", line 23, in <cell line: 0>
    xla_out = tf.function(model_5415, jit_compile=True)(*_args)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorflow/python/util/traceback_utils.py", line 167, in error_handler
    raise e.with_traceback(filtered_tb) from None
  File "/usr/local/lib/python3.12/dist-packages/tensorflow/python/eager/execute.py", line 59, in quick_execute
    except TypeError as e:
tensorflow.python.framework.errors_impl.InvalidArgumentError: Dynamic inferencing on custom call goo.gle/debugstr  
name: "Qr.4"
opcode: "custom-call"
shape {
  element_type: F16
  dimensions: 3
  dimensions: 32
  dimensions: 32
  layout {
    minor_to_major: 2
    minor_to_major: 1
    minor_to_major: 0
    tail_padding_alignment_in_elements: 1
  }
  is_dynamic_dimension: false
  is_dynamic_dimension: false
  is_dynamic_dimension: false
}
metadata {
  op_type: "Qr"
  op_name: "Qr"
  source_file: "/usr/local/lib/python3.12/dist-packages/tensorflow/python/framework/ops.py"
  source_line: 1221
}
custom_call_target: "ProductOfElementaryHouseholderReflectors"
id: 5
operand_ids: 3
operand_ids: 4
frontend_attributes {
}
custom_call_api_version: API_VERSION_ORIGINAL
 is not supported
	while evaluating input 1 of Range operator as a compile-time constant.

Stack trace for op definition: 
File "<frozen runpy>", line 198, in _run_module_as_main
File "<frozen runpy>", line 88, in _run_code
File "/usr/local/lib/python3.12/dist-packages/colab_kernel_launcher.py", line 37, in <module>
File "/usr/local/lib/python3.12/dist-packages/traitlets/config/application.py", line 992, in launch_instance
File "/usr/local/lib/python3.12/dist-packages/ipykernel/kernelapp.py", line 712, in start
File "/usr/local/lib/python3.12/dist-packages/tornado/platform/asyncio.py", line 211, in start
File "/usr/lib/python3.12/asyncio/base_events.py", line 645, in run_forever
File "/usr/lib/python3.12/asyncio/base_events.py", line 1999, in _run_once
File "/usr/lib/python3.12/asyncio/events.py", line 88, in _run
File "/usr/local/lib/python3.12/dist-packages/ipykernel/kernelbase.py", line 510, in dispatch_queue
File "/usr/local/lib/python3.12/dist-packages/ipykernel/kernelbase.py", line 499, in process_one
File "/usr/local/lib/python3.12/dist-packages/ipykernel/kernelbase.py", line 406, in dispatch_shell
File "/usr/local/lib/python3.12/dist-packages/ipykernel/kernelbase.py", line 730, in execute_request
File "/usr/local/lib/python3.12/dist-packages/ipykernel/ipkernel.py", line 383, in do_execute
File "/usr/local/lib/python3.12/dist-packages/ipykernel/zmqshell.py", line 528, in run_cell
File "/usr/local/lib/python3.12/dist-packages/IPython/core/interactiveshell.py", line 2975, in run_cell
File "/usr/local/lib/python3.12/dist-packages/IPython/core/interactiveshell.py", line 3030, in _run_cell
File "/usr/local/lib/python3.12/dist-packages/IPython/core/async_helpers.py", line 78, in _pseudo_sync_runner
File "/usr/local/lib/python3.12/dist-packages/IPython/core/interactiveshell.py", line 3257, in run_cell_async
File "/usr/local/lib/python3.12/dist-packages/IPython/core/interactiveshell.py", line 3473, in run_ast_nodes
File "/usr/local/lib/python3.12/dist-packages/IPython/core/interactiveshell.py", line 3553, in run_code
File "/tmp/ipykernel_2124/288000946.py", line 23, in <cell line: 0>
File "/tmp/ipykernel_2124/1488513231.py", line 9, in model_5415

	 [[{{node SequenceMask/Range}}]]
	tf2xla conversion failed while converting __inference_model_5415_900[_XlaMustCompile=true,config_proto=13561319589895757934,executor_type=11160318154034397263]. Run with TF_DUMP_GRAPH_PREFIX=/path/to/dump/dir and --vmodule=xla_compiler=2 to obtain a dump of the compiled functions. [Op:__inference_model_5415_900]

Version

Tensorflow: 2.21.0