End-to-end Learning in Ray-tracing channel

[Mio-TW]

Dear Sionna Team,

I want to run end-to-end learning in a ray-tracing channel.
So I combined the tutorials from Autoencoder and Ray-tracing, and used the conventional training in my codes (not reinforcement learning case).
The channel is generated using the following codes in here

# Init CIR generator
cir_generator = CIRGenerator(a,
                             tau,
                             num_tx)

# Initialises a channel model that can be directly used by OFDMChannel layer
channel_model = CIRDataset(cir_generator,
                           batch_size,
                           num_rx,
                           num_rx_ant,
                           num_tx,
                           num_tx_ant,
                           75,
                           num_time_steps)

# Configure the actual channel
self._channel = OFDMChannel(
                     self._channel_model,
                     self._pusch_transmitter.resource_grid,
                     normalize_channel=True,
                     return_channel=True)

But I suffer the following error.
It seems that the gradient cannot be backpropagated through the "ofdm channel".

So I have the following questions:

1. Why can't the gradient pass through the OFDM channel?
   Haven't we already calculated the channel response using 'a' and 'tau'? (Equivalent to knowing 'H' in y=Hx+n)
   Just like passing through an AWGN channel, shouldn't it be enough to pass the gradient through?
   I'm not sure why this error would happen.
   Please tell me if my understanding is correct.
   If so, please tell me how to change the code of OFDM channel so that it can pass gradient.

2. If I use the "tf.stop_gradient" suggested by the error message, the trainable constellation parameters on the transmitter side will not be updated, which is not the result I want.

Thank you so much.

File ~/.local/lib/python3.8/site-packages/tensorflow/python/ops/gradients_util.py:631, in _GradientsHelper(ys, xs, grad_ys, name, colocate_gradients_with_ops, gate_gradients, aggregation_method, stop_gradients, unconnected_gradients, src_graph)
    629       grad_fn = func_call.python_grad_func
    630     else:
--> 631       raise LookupError(
    632           "No gradient defined for operation"
    633           f"'{op.name}' (op type: {op.type}). "
    634           "In general every operation must have an associated "
    635           "`@tf.RegisterGradient` for correct autodiff, which this "
    636           "op is lacking. If you want to pretend this "
    637           "operation is a constant in your program, you may insert "
    638           "`tf.stop_gradient`. This can be useful to silence the "
    639           "error in cases where you know gradients are not needed, "
    640           "e.g. the forward pass of tf.custom_gradient. "
    641           "Please see more details in "
    642           "https://www.tensorflow.org/api_docs/python/tf/custom_gradient.")  # pylint: disable=line-too-long
    643 if loop_state:
    644   loop_state.EnterGradWhileContext(op, before=False)

No gradient defined for operation'ofdm_channel_1/IteratorGetNext' (op type: IteratorGetNext). 
In general every operation must have an associated `@tf.RegisterGradient` for correct autodiff, which this op is lacking. 
If you want to pretend this operation is a constant in your program, you may insert `tf.stop_gradient`. 
This can be useful to silence the error in cases where you know gradients are not needed, e.g. the forward pass of tf.custom_gradient.
Please see more details in https://www.tensorflow.org/api_docs/python/tf/custom_gradient.

[faycalaa]

Hi @Mio-TW,

Could you please share a minimal code snippet that reproduces the error?

[Mio-TW]

Hi @faycalaa,
please kindly help check it.
Thank you so much.

[jhoydis]

Hi @Mio-TW,

Have you tried deactivating XLA, i.e., @tf.function(jit_compile=False), or running the training in Eager mode, i.e., remove the tf.function decorator completely? I think that the CIR generator does not work in XLA.

[Mio-TW]

Hi @jhoydis,

I added the command sionna.config.xla_compat=False at the beginning of the code.
Also the @tf.function(jit_compile=False) before the calling function of CIRGenerator.
But still got the same error.

[jhoydis]

Did you change the decorator of the call method of the class E2ESystemConventionalTraining?

[Mio-TW]

Hi @jhoydis,
Yes I also changed it into False.