llr = self._neural_receiver([y, no])
# [batch_size, num_rx, num_streams_per_rx, num_ofdm_symbols, fft_size, num_bits_per_symbol]
llr = insert_dims(llr, 2, 1)
# [batch_size, num_rx, num_streams_per_rx, num_data_symbols, num_bits_per_symbol]
llr = self._rg_demapper(llr)
# [batch_size, num_tx, 1, num_data_symbols*num_bit_per_symbols]
llr = flatten_last_dims(llr, 2)
# Remove stream dimension, NOTE: does not support
# multiple-streams
# per user; conceptually the neural receiver does, but would
# require modified reshapes
llr = tf.squeeze(llr, axis=-2)
[batch size, num ofdm symbols, num subcarriers, num_bits_per_symbol]
How to change the code if num_streams_per_rx is not equal to 1