pad_right for tensorflow

pad_right_for_tensorflow.py

import tensorflow as tf
from tensorflow import Tensor, TensorShape
from typing import Union


def pad_right(
    x: Tensor, target_shape: TensorShape, value: Union[int, float] = 0
) -> Tensor:
    """Creates a tensor of shape `target_shape` by padding it with `value` on the right.
    
    Args:
        x: Any tensor
        target_shape: Shape of the padded x; must be at least as large as the
            shape of x in all dimensions

    Returns:
        A tensor like x, but padded with zeros
    """
    current_shape = tf.shape(x)
    right_padding = tf.expand_dims(
        tf.convert_to_tensor(target_shape - current_shape), -1
    )
    padding = tf.concat([tf.zeros_like(right_padding), right_padding], -1)
    return tf.pad(x, padding, "CONSTANT", constant_values=value)
    
    
##############################################################################


import pytest
import tensorflow as tf
import numpy as np
import rasa.utils.tensorflow.layers_utils as layers_utils


def test_pad_right():
    x = tf.ones([3, 2])
    x_padded = layers_utils.pad_right(x, [5, 7])
    assert np.all(tf.shape(x_padded).numpy() == [5, 7])
    assert np.all(
        x_padded.numpy()
        == [
            [1, 1, 0, 0, 0, 0, 0],
            [1, 1, 0, 0, 0, 0, 0],
            [1, 1, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
            [0, 0, 0, 0, 0, 0, 0],
        ]
    )