XLA compilation fails with InvalidArgumentError for tf.linalg.svd on zero-sized dimensions

[jasminetrail]

XLA compilation fails when tf.linalg.svd is applied to the output of a tf.keras.layers.Embedding layer with output_dim=0, while the same program executes successfully in eager mode.

InvalidArgumentError: Slice dim size 1 greater than dynamic slice dimension: 0. [[{{node Svd}}]]

This indicates an inconsistency between eager execution and XLA compilation when handling this computation. Rather than producing the same behavior as eager execution or reporting a clear validation error for an unsupported case, XLA terminates with an internal compilation error originating from the Svd operation. Here is the gist

Minimal Code

import traceback
import tensorflow as tf

tf.random.set_seed(8761)
t23420 = tf.complex(tf.random.normal([1, 4, 16, 16, 32], dtype=tf.float32), tf.random.normal([1, 4, 16, 16, 32], dtype=tf.float32))

_m_t23422 = tf.keras.layers.Embedding(8, 0)

def model_8761(t23420):
    t23422 = _m_t23422(t23420)
    t23424 = tf.linalg.svd(t23422, full_matrices=False, compute_uv=False)
    return t23424

_args = (t23420,)

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

print("\n=== XLA ===")
try:
    xla_out = tf.function(model_8761, 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_1541/2978194018.py", line 25, in <cell line: 0>
    xla_out = tf.function(model_8761, 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: Slice dim size 1 greater than dynamic slice dimension: 0.
	 [[{{node Svd}}]]
	tf2xla conversion failed while converting __inference_model_8761_572[_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_8761_572]

Versions

Tensorflow: 2.21.0