mirror of
https://github.com/twitter/the-algorithm.git
synced 2024-10-03 20:25:11 +02:00
61 lines
1.8 KiB
Python
61 lines
1.8 KiB
Python
import numpy as np
|
|
import tensorflow.compat.v1 as tf
|
|
|
|
|
|
def create_sparse_tensor(
|
|
batch_size: int,
|
|
input_size: int,
|
|
num_values: int,
|
|
dtype: tf.DType = tf.float32,
|
|
) -> tf.SparseTensor:
|
|
"""
|
|
Creates a sparse tensor with the given batch size, input size, and number of values.
|
|
|
|
Args:
|
|
batch_size (int): The batch size of the sparse tensor.
|
|
input_size (int): The input size of the sparse tensor.
|
|
num_values (int): The number of values in the sparse tensor.
|
|
dtype (tf.DType): The dtype of the sparse tensor.
|
|
|
|
Returns:
|
|
A sparse tensor with the given batch size, input size, and number of values.
|
|
"""
|
|
random_indices = np.sort(
|
|
np.random.randint(batch_size * input_size, size=num_values)
|
|
)
|
|
test_indices_i = random_indices // input_size
|
|
test_indices_j = random_indices % input_size
|
|
test_indices = np.stack([test_indices_i, test_indices_j], axis=1)
|
|
test_values = np.random.random(num_values).astype(dtype.as_numpy_dtype)
|
|
|
|
return tf.SparseTensor(
|
|
indices=tf.constant(test_indices),
|
|
values=tf.constant(test_values),
|
|
dense_shape=(batch_size, input_size),
|
|
)
|
|
|
|
|
|
def create_reference_input(
|
|
sparse_input: tf.SparseTensor, use_binary_values: bool
|
|
) -> tf.SparseTensor:
|
|
"""
|
|
Creates a reference input for the sparse input.
|
|
|
|
Args:
|
|
sparse_input (tf.SparseTensor): The sparse input.
|
|
use_binary_values (bool): Whether to use binary values.
|
|
|
|
Returns:
|
|
A reference input for the sparse input.
|
|
"""
|
|
|
|
if use_binary_values:
|
|
sp_a = tf.SparseTensor(
|
|
indices=sparse_input.indices,
|
|
values=tf.ones_like(sparse_input.values),
|
|
dense_shape=sparse_input.dense_shape,
|
|
)
|
|
else:
|
|
sp_a = sparse_input
|
|
return sp_a
|