How to workaround the SVD issue with XLA

[pustovale]

Hi all,

I implemented SVD-based precoding in Sionna, but encountered the issue which I'm struggling to resolve.
The code adds SVD-precoding to the code of "MIMO OFDM Transmissions over the CDL Channel Model" using the Tensorflow function:
s, u, v = tf.linalg.svd(h)

However, when running the code I encountered the error:
No registered 'Svd' OpKernel for XLA_CPU_JIT

I found that when disable XLA optimization for all the code

@tf.function(jit_compile=False)
def call(self, batch_size, ebno_db):

everything works well.
However, this way the code performance is degraded. I tried to workaround the SVD function with the same decorator @tf.function(jit_compile=False) but it does not help to resolve the error if the rest of the code is run with enabled optimization.

Thus I have a question is there any other way to disable optimization for SVD while keeping the rest of the code with XLA.

[jhoydis]

Hi!

You can in principle put decorators around any part of the model you want. But you cannot put an XLA decorator around the entire model and then selectively exclude parts of the graph from it. So you would need to decorate specific components, e.g., the channel model and the receiver.

You could also try to implement SVD in a different manner. If an eigen decomposition would be sufficient, you could use tf.linalg.eigh which works with XLA. I am guessing that you might only need either u or v. You could then probably do something like this:

_, u = tf.linalg.eigh(tf.matmul(h,h,adjoint_b=True))
_, v = tf.linalg.eigh(tf.matmul(h,h,adjoint_a=True))

[pustovale]

Thanks! The suggestion with tf.linalg.eigh works well. For whom it may concern, don't forget to reverse the order of the columns in v with v = tf.reverse(v,axis=[-1]) as eigen values in eigh are sorted in the descend order, opposite to the svd.