def ShiftRight(x, **unused_kwargs):
  """Layer to shift the tensor to the right by padding on axis 1."""
  if not isinstance(x, (list, tuple)):  # non-chunked inputs
    pad_widths = [(0, 0), (1, 0)]
    padded = np.pad(x, pad_widths, mode='constant')
    return padded[:, :-1]
  # Handling chunked inputs. Recall that the list of chunks represents a big
  # sequence (the concatenation of the chunks). We want to shift that sequence,
  # so we put a 0 in the beginning of the first chunk and the last element of
  # that chunk is used as the new first element of the next chunk, and so on.
  padded = []
  last_value = np.zeros_like(x[0][:, -1])
  for chunk in x:
    padded_chunk = np.concatenate([last_value[:, np.newaxis], chunk], axis=1)
    last_value = chunk[:, -1]
    padded.append(padded_chunk[:, :-1])
  return padded
Layer to shift the tensor to the right by padding on axis 1.