help converting baseline from autoencoder to multiple sample/symbol

[nbecker]

I've been working on converting the example Autoencoder to multiple sample/symbol, using the new upsampling/downsampling filters. I've gotten stuck on an error in TF that occurs during the tracing when run in graph mode. The error message provides no information to me, but when run in eager mode there is no error. I'm wondering if anyone maybe has more experience in TF and might have some ideas on what the problem might be or on how to debug.

I've placed a minimal sample code here:
https://gist.github.com/nbecker/d1bf599b9fb8e022a292017aab707e3f

The error is

TypeError: Exception encountered when calling layer "baseline_6" (type Baseline).

in user code:

    File "/home/nbecker/ml-phy/baseline_2sps.py", line 119, in call  *
        llr = self._demapper([y, no])
    File "/home/nbecker/.local/lib/python3.10/site-packages/keras/utils/traceback_utils.py", line 67, in error_handler  **
        raise e.with_traceback(filtered_tb) from None

    TypeError: Exception encountered when calling layer "demapper_6" (type Demapper).
    
    in user code:
    
        File "/home/nbecker/.local/lib/python3.10/site-packages/sionna/mapping.py", line 818, in call  *
            return super().call([y, null_prior, no])
        File "/home/nbecker/.local/lib/python3.10/site-packages/sionna/mapping.py", line 702, in call  *
            out_shape = tf.concat([tf.shape(y)[:-1],
    
        TypeError: unsupported operand type(s) for *: 'NoneType' and 'int'

[faycalaa]

Thank you for sharing your code. It helped us unveil a bug in the function sionna.utils.convolution(). In the current release (v0.9), the output shape of this function can be variable for dimensions other than the first one, causing the error you observe.
We will fix with this bug in the next Sionna release.

[nbecker]

As a workaround I have an alternative implementation of APP demapper, based on code I've used for years.
First enumerate all symbols which have 0's and 1's in each bit position:

def gen_ind (bits):
    size = 1<<bits
    all_syms = np.arange (size)
    all_bits = np.arange (bits)
    bit_mask = (all_syms[:,np.newaxis] >> all_bits) & 1

    ind1 = np.array ([np.nonzero (bit_mask[...,i])[0] for i in range (bits)])
    ind0 = np.array ([np.nonzero (np.logical_not (bit_mask[...,i]))[0] for i in range (bits)])
    return ind1, ind0

ind1, ind0 = gen_ind (bits_per_sym)

Now in numpy and in eager mode we can use these indices to directly index (as in the code commented out below). Unfortunately this isn't working in graph mode, and we have to resort to the un-commented code which enumerates all indices and uses tf.gather_nd.

from tensorflow.python.ops.numpy_ops import np_config
np_config.enable_numpy_behavior()

class demapper ():
    def __init__ (self, constellation, M, n_symbols):
        self.constellation = constellation
        all_indices = np.array([[[[b,i,j] for j in range (M)] for i in range (n_symbols)]  for b in range (training_batch_size.numpy())] )
        self.indices_0 = all_indices[:,:,ind0]
        self.indices_1 = all_indices[:,:,ind1]
    def __call__ (self, inputs, var):
        #inputs[batches, N]
        dist = tf.math.square (tf.abs (inputs[:,:,np.newaxis] - self.constellation.points))
        ll = -dist / var[:,:,None]
        # Unforturnately this direct addressing doesn't work in graph mode
        # ld0 = ll[:,:,ind0]
        # ld1 = ll[:,:,ind1]
        # Instead we have to enumerate all the indices and use gather_nd
        ld0 = tf.gather_nd (ll, self.indices_0)
        ld1 = tf.gather_nd (ll, self.indices_1)
        llr = tf.math.reduce_logsumexp (ld0, axis=-1) - tf.math.reduce_logsumexp (ld1, axis=-1)
        # llr has shape [batch, n_symbols, bits_per_sym]
        llr_flat = tf.reshape (llr[:,:,::-1], (llr.shape[0],N))
        return -llr_flat        # wants llrs inverted