XLA compilation fails for tf.linalg.matrix_transpose on keras.layers.ReversibleEmbedding output with zero-sized embedding dimension

[jasminetrail]

tf.linalg.matrix_transpose behaves inconsistently between eager execution and XLA when applied to the output of tf.keras.layers.ReversibleEmbedding configured with input_dim=0 and output_dim=0.

In eager mode, ReversibleEmbedding produces a tensor with shape (3, 8, 32, 32, 32, 0), and tf.linalg.matrix_transpose executes successfully. When the same function is compiled with tf.function(jit_compile=True), XLA compilation fails with:

InvalidArgumentError: transpose expects a vector of size 1. But input(1) is a vector of size 6

The eager and XLA executions observe the same tensor shape ((3, 8, 32, 32, 32, 0)), but XLA appears to infer an incorrect rank for the input to tf.linalg.matrix_transpose, causing the transpose permutation validation to fail during tf2xla conversion. Here is the gist

Minimal code

import traceback
import tensorflow as tf

tf.random.set_seed(3130)
t6375 = tf.random.normal([3, 8, 32, 32, 32], dtype=tf.float64)

_m_t6376 = tf.keras.layers.ReversibleEmbedding(0, 0)

def model_3130(t6375):
    t6376 = _m_t6376(t6375)
    print("t6376.shape", t6376.shape)
    t6379 = tf.linalg.matrix_transpose(t6376)
    return t6379

_args = (t6375,)

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

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

Error logs

=== EAGER ===
t6376.shape (3, 8, 32, 32, 32, 0)
Eager Output Succeeded

=== XLA ===
t6376.shape (3, 8, 32, 32, 32, 0)
Traceback (most recent call last):
  File "/tmp/ipykernel_1541/1942597633.py", line 26, in <cell line: 0>
    xla_out = tf.function(model_3130, 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: transpose expects a vector of size 1. But input(1) is a vector of size 6

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_1541/1942597633.py", line 26, in <cell line: 0>
File "/tmp/ipykernel_1541/1942597633.py", line 12, in model_3130

	 [[{{node matrix_transpose/transpose}}]]
	tf2xla conversion failed while converting __inference_model_3130_264[_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_3130_264]

versions

Tensorflow: 2.21.0