Слои trax tl.Relu и tl.ShiftRight вложены в Serial Combinator

Question

Я пытаюсь построить модель внимания, но слои Relu и ShiftRight по умолчанию вложены в Serial Combinator. Это также дает мне ошибки в обучении.

      layer_block = tl.Serial(
    tl.Relu(),
    tl.LayerNorm(), )

x = np.array([[-2, -1, 0, 1, 2],
              [-20, -10, 0, 10, 20]]).astype(np.float32) 

layer_block.init(shapes.signature(x)) y = layer_block(x)

print(f'layer_block: {layer_block}')

Выход

      layer_block: Serial[
  Serial[
    Relu
  ]
  LayerNorm
]

Ожидаемый результат

      layer_block: Serial[
  Relu
  LayerNorm
]

Та же проблема возникает с tl.ShiftRight ()

Приведенный выше код взят из официальной документации Пример 5.

заранее спасибо

Answer 1

Я не смог найти точного решения указанной выше проблемы, но вы можете создать собственную функцию с помощью tl.Fn() и добавить в код функции Relu и нееShiftRight .

      def _zero_pad(x, pad, axis):
    """Helper for jnp.pad with 0s for single-axis case."""
    pad_widths = [(0, 0)] * len(x.shape)
    pad_widths[axis] = pad  # Padding on axis.
    
    return jnp.pad(x, pad_widths, mode='constant')


def f(x):
    if mode == 'predict':
        return x
    padded = _zero_pad(x, (n_positions, 0), 1)
    return padded[:, :-n_positions]

# set ShiftRight parameters as global 
n_positions = 1
mode='train'

layer_block = tl.Serial(
    tl.Fn('Relu', lambda x: jnp.where(x <= 0, jnp.zeros_like(x), x)),
    tl.LayerNorm(),
    tl.Fn(f'ShiftRight({n_positions})', f)
)


x = np.array([[-2, -1, 0, 1, 2],
              [-20, -10, 0, 10, 20]]).astype(np.float32)
layer_block.init(shapes.signature(x))
y = layer_block(x)


print(f'layer_block: {layer_block}')

Выход

      layer_block: Serial[
  Relu
  LayerNorm
  ShiftRight(1)
]