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]))