def einsum_batched_matrix_high_rank_vector_mul(module):
module.einsum_batched_matrix_high_rank_vector_mul(
tf_utils.ndarange([BATCH_DIM, 2, 6]),
tf_utils.ndarange([BATCH_DIM, 5, 3, 6]))
def einsum_mul(module):
module.einsum_mul(tf_utils.ndarange([VECTOR_DIM]),
tf_utils.ndarange([VECTOR_DIM]))
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_rank_diag(module):
module.einsum_dynamic_rank_diag(
tf_utils.ndarange([BATCH_DIM, BATCH_DIM, LEFT_DIM, RIGHT_DIM]))
def einsum_identity(module):
module.einsum_identity(tf_utils.ndarange([VECTOR_DIM]))
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 einsum_dynamic_rank_identity(module):
module.einsum_dynamic_rank_identity(
tf_utils.ndarange([BATCH_DIM, LEFT_DIM, RIGHT_DIM]))
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]))
def downsample_nearest_neighbor(module):
img = tf_utils.ndarange([1, 52, 37, 1], dtype=np.int32)
module.downsample_nearest_neighbor(img)
def batched_feature_unpadded(module):
i = tf_utils.ndarange([2, 4, 5, 2])
k = tf_utils.ndarange([2, 2, 2, 3])
module.conv2d_2452x2423_valid(i, k)
def batched_feature_padded_same_stride_2(module):
i = tf_utils.ndarange([2, 4, 5, 2])
k = tf_utils.ndarange([2, 4, 2, 3])
module.conv2d_2452x2423_same_stride_2(i, k)
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 gather_axis1_batch1(module):
indices = np.array([[2], [3], [0], [1]], dtype=np.int32)
params = tf_utils.ndarange([4, 7, 8, 2])
module.gather_axis1_batch1(params, indices)
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 upsample_nearest_neighbor(module):
img = tf_utils.ndarange([1, 8, 7, 1], dtype=np.int32)
module.upsample_nearest_neighbor(img)
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 einsum_implicit_transpose(module): module.einsum_implicit_transpose(tf_utils.ndarange([LEFT_DIM, RIGHT_DIM]))
def einsum_dynamic_dim_identity(module):
module.einsum_dynamic_dim_identity(
tf_utils.ndarange([LEFT_DIM, RIGHT_DIM]))
def einsum_explicit_trace(module): module.einsum_explicit_trace(tf_utils.ndarange([LEFT_DIM, RIGHT_DIM]))
def einsum_dynamic_dim_transpose(module):
module.einsum_dynamic_dim_transpose(
tf_utils.ndarange([BATCH_DIM, LEFT_DIM, RIGHT_DIM]))
def einsum_diag(module): module.einsum_diag(tf_utils.ndarange([LEFT_DIM, RIGHT_DIM]))
def einsum_outer_product(module):
module.einsum_outer_product(tf_utils.ndarange([VECTOR_DIM]),
tf_utils.ndarange([VECTOR_DIM]))
def einsum_sum_axis_1(module): module.einsum_sum_axis_1(tf_utils.ndarange([LEFT_DIM, RIGHT_DIM]))
def einsum_sum(module):
module.einsum_sum(tf_utils.ndarange([VECTOR_DIM]))
def einsum_broadcast_singleton_dimension(module):
module.einsum_broadcast_singleton_dimension(
tf_utils.ndarange([1, LEFT_DIM, INNER_DIM]),
tf_utils.ndarange([BATCH_DIM, INNER_DIM, RIGHT_DIM]))
def einsum_explicit_inner_product(module):
module.einsum_explicit_inner_product(tf_utils.ndarange([VECTOR_DIM]),
tf_utils.ndarange([VECTOR_DIM]))
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=BINARY_SIGNATURE_SHAPES, signature_dtypes=[tf.bool]),
"maximum":
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=BINARY_SIGNATURE_SHAPES,
signature_dtypes=[tf.float32, tf.int32]),
"minimum":
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=BINARY_SIGNATURE_SHAPES,
signature_dtypes=[tf.float32, tf.int32]),
"mod":
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=BINARY_SIGNATURE_SHAPES,
signature_dtypes=[tf.float32, tf.int32],
input_generators={
"positive_ndarange": lambda *args: tf_utils.ndarange(*args) + 1
}),
"multiply":
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=BINARY_SIGNATURE_SHAPES,
signature_dtypes=[tf.float32, tf.int32, tf.complex64]),
"multiply_no_nan":
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=BINARY_SIGNATURE_SHAPES,
signature_dtypes=[tf.float32, tf.complex64]),
"ndtri":
tf_test_utils.unit_test_specs_from_signatures(
signature_shapes=UNARY_SIGNATURE_SHAPES,
signature_dtypes=[tf.float32]),
"negative":
tf_test_utils.unit_test_specs_from_signatures(