def batched_feature_padded_same_stride_1_output_1(module):
i = tf_utils.ndarange([2, 4, 5, 4])
k = tf_utils.ndarange([2, 4, 4, 1])
module.conv2d_2453x2441_same_stride_1(i, k)
def batched_feature_unpadded(module):
i = tf_utils.ndarange([2, 4, 5, 2])
k = tf_utils.ndarange([2, 2, 2, 3])
module.conv2d_2452x2223_valid(i, k)
def downsample_nearest_neighbor(module): img = tf_utils.ndarange([1, 52, 37, 1], dtype=np.int32) module.downsample_nearest_neighbor(img)
def feature_reduce(module):
i = tf_utils.ndarange([1, 4, 5, 2])
k = np.ones([3, 2, 2, 1], dtype=np.float32)
module.conv2d_1452x3221_same(i, k)
def feature_mix(module):
i = tf_utils.ndarange([1, 4, 5, 2])
k = tf_utils.ndarange([1, 1, 2, 2])
module.conv2d_1452x1122_same(i, k)
def gather_axis2_batch1(module): indices = np.array([[2], [3], [0], [1]], dtype=np.int32) params = tf_utils.ndarange([4, 7, 8, 2]) module.gather_axis2_batch1(params, indices)
def asymmetric_kernel(module):
i = tf_utils.ndarange([1, 4, 5, 1])
k = np.array([[1, 4, 2], [-2, 0, 1]],
dtype=np.float32).reshape(2, 3, 1, 1)
module.conv2d_1451x2311_valid(i, k)
def einsum_explicit_inner_product(module):
module.einsum_explicit_inner_product(tf_utils.ndarange([VECTOR_DIM]),
tf_utils.ndarange([VECTOR_DIM]))
def transposed(module):
kernel = tf_utils.ndarange([1, 16, 16, 32])
img = tf_utils.ndarange([1, 1, 32, 32])
module.conv2d_transpose_same(kernel, img)
def einsum_sum(module): module.einsum_sum(tf_utils.ndarange([VECTOR_DIM]))
def einsum_mul(module):
module.einsum_mul(tf_utils.ndarange([VECTOR_DIM]),
tf_utils.ndarange([VECTOR_DIM]))
def einsum_identity(module): module.einsum_identity(tf_utils.ndarange([VECTOR_DIM]))
def einsum_outer_product(module):
module.einsum_outer_product(tf_utils.ndarange([VECTOR_DIM]),
tf_utils.ndarange([VECTOR_DIM]))
def batched_feature_unpadded_same_stride_2(module):
i = tf_utils.ndarange([2, 4, 5, 2])
k = tf_utils.ndarange([2, 4, 2, 3])
module.conv2d_2452x2423_valid_stride_2(i, k)
def gather_axis0_scalar(module): indices = np.array(2, dtype=np.int32) params = tf_utils.ndarange([4, 8]) module.gather_axis0_scalar(params, indices)
def einsum_dynamic_dim_sum(module):
module.einsum_dynamic_dim_sum(
tf_utils.ndarange([BATCH_DIM, BATCH_DIM, LEFT_DIM, RIGHT_DIM]))
def gather_axis1_batch0(module): indices = np.array([2, 3], dtype=np.int32) params = tf_utils.ndarange([4, 7, 8]) module.gather_axis1_batch0(params, indices)
def einsum_dynamic_dim_matmul(module):
module.einsum_dynamic_dim_matmul(
tf_utils.ndarange([LEFT_DIM, INNER_DIM]),
tf_utils.ndarange([INNER_DIM, RIGHT_DIM]))
def id_batch_size_2(module):
i = tf_utils.ndarange([2, 4, 5, 1])
k = np.ones([1, 1, 1, 1], dtype=np.float32)
module.conv2d_2451x1111_valid(i, k)
def einsum_dynamic_dim_lhs_batch(module):
module.einsum_dynamic_dim_lhs_batch(
tf_utils.ndarange([BATCH_DIM, LEFT_DIM, INNER_DIM]),
tf_utils.ndarange([INNER_DIM, RIGHT_DIM]))
def batched_padding(module):
i = tf_utils.ndarange([2, 4, 5, 1])
k = np.array([[1, 4, 2], [-2, 0, 1]],
dtype=np.float32).reshape(2, 3, 1, 1)
module.conv2d_2451x2311_same(i, k)
def einsum_dynamic_rank_split_heads(module):
module.einsum_dynamic_rank_split_heads(
tf_utils.ndarange([BATCH_DIM, BATCH_DIM, 8, 6]),
tf_utils.ndarange([12, 6, 4]))
def feature_inflate(module):
i = tf_utils.ndarange([1, 4, 5, 1])
k = tf_utils.ndarange([1, 1, 1, 2])
module.conv2d_1451x1112_same(i, k)
def einsum_dynamic_dim_identity(module):
module.einsum_dynamic_dim_identity(
tf_utils.ndarange([LEFT_DIM, RIGHT_DIM]))
def feature_padded(module):
i = tf_utils.ndarange([1, 4, 5, 2])
k = tf_utils.ndarange([2, 2, 2, 3])
module.conv2d_1452x2223_same(i, k)
def einsum_dynamic_rank_identity(module):
module.einsum_dynamic_rank_identity(
tf_utils.ndarange([BATCH_DIM, LEFT_DIM, RIGHT_DIM]))
def predict(module):
inputs = tf_utils.ndarange(INPUT_SHAPE)
module.predict(inputs)
def predict(module):
inputs = tf_utils.ndarange(INPUT_SHAPE)
module.predict(inputs, rtol=1e-5, atol=1e-5)
def upsample_nearest_neighbor(module): img = tf_utils.ndarange([1, 8, 7, 1], dtype=np.int32) module.upsample_nearest_neighbor(img)
def slice_first_element_with_from_tensor_high_rank(module):
module.slice_first_element_with_from_tensor_high_rank(
tf_utils.ndarange([STATIC_SIZE, STATIC_SIZE]))