Как перебрать все элементы Tensor и получить индекс

Question

У меня есть модель Keras Sequential с настроенными плотными слоями, и я хочу перебрать тензор ядра и смещения каждого слоя, имея возможность получить доступ к значению элемента и индексу, где он находится. У меня есть следующий код:

l = []
for i in range(len(nn.layers)):
    t1, t2 = nn.layers[i].kernel, nn.layers[i].bias
    for j in range(t1.shape[0]):
        for k in range(t1[j].shape[0]):
            l.append((t1[j][k].numpy(), i, j, k))
    for j in range(t2.shape[0]):
        l.append((t2[j], i))

Где переменная nn является последовательной моделью с определенным количеством слоев, каждый из которых имеет количество узлов. Это работает для моих целей, но не очень эффективно. Есть ли какой-либо (желательно более TensorFlow-подобный способ) достижения того же результата? Кстати, я использую TensorFlow 2.0.