# [batch size, num ofdm symbols, num subcarriers, num_bits_per_symbol] 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) How to change the code if num_streams_per_rx is not equal to 1