Skip to content

Commit

Permalink
Adds performance enhancements for sparse embedding lookups.
Browse files Browse the repository at this point in the history
  • Loading branch information
@philipphack
philipphack committed Feb 16, 2023
1 parent c393302 commit 5165f07
Show file tree
Hide file tree
Showing 8 changed files with 40 additions and 6 deletions.
4 changes: 3 additions & 1 deletion tensorflow/python/eager/pywrap_gradient_exclusions.cc
View file
Expand Up @@ -235,6 +235,7 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(
{"QueueDequeueMany"},
{"QueueDequeueUpTo"},
{"QueueSize"},
{"RaggedFillEmptyRows"},
{"RaggedRange"},
{"RandomCrop"},
{"RandomIndexShuffle"},
Expand Down Expand Up @@ -428,7 +429,7 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedInputIndices(

absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
const tensorflow::string &op_name) {
static std::array<OpIndexInfo, 483> a = {{
static std::array<OpIndexInfo, 484> a = {{
{"Abs"},
{"AccumulateNV2"},
{"Acos"},
Expand Down Expand Up @@ -689,6 +690,7 @@ absl::optional<tensorflow::gtl::FlatSet<int>> OpGradientUnusedOutputIndices(
{"QueueSize"},
{"RFFT"},
{"RFFT2D"},
{"RaggedFillEmptyRows", 3, {0, 1, 2}},
{"RaggedGather"},
{"RaggedRange"},
{"RaggedTensorToSparse"},
Expand Down
2 changes: 1 addition & 1 deletion tensorflow/python/ops/sparse_ops.py
View file
Expand Up @@ -1067,7 +1067,7 @@ def sparse_split_v2(sp_input=None,
>>> indices = [[0, 2], [0, 4], [0, 5], [1, 0], [1, 1]]
>>> values = [1, 2, 3, 4, 5]
>>> t = tf.sparse.SparseTensor(indices=indices, values=values,
dense_shape=[2, 7])
... dense_shape=[2, 7])
>>> tf.sparse.to_dense(t)
<tf.Tensor: shape=(2, 7), dtype=int32, numpy=
array([[0, 0, 1, 0, 2, 3, 0],
Expand Down
4 changes: 2 additions & 2 deletions tensorflow/tools/api/golden/v1/tensorflow.nn.pbtxt
View file
Expand Up @@ -202,7 +202,7 @@ tf_module {
}
member_method {
name: "embedding_lookup_sparse"
argspec: "args=[\'params\', \'sp_ids\', \'sp_weights\', \'partition_strategy\', \'name\', \'combiner\', \'max_norm\'], varargs=None, keywords=None, defaults=[\'mod\', \'None\', \'None\', \'None\'], "
argspec: "args=[\'params\', \'sp_ids\', \'sp_weights\', \'partition_strategy\', \'name\', \'combiner\', \'max_norm\', \'allow_fast_lookup\'], varargs=None, keywords=None, defaults=[\'mod\', \'None\', \'None\', \'None\', \'False\'], "
}
member_method {
name: "erosion2d"
Expand Down Expand Up @@ -338,7 +338,7 @@ tf_module {
}
member_method {
name: "safe_embedding_lookup_sparse"
argspec: "args=[\'embedding_weights\', \'sparse_ids\', \'sparse_weights\', \'combiner\', \'default_id\', \'name\', \'partition_strategy\', \'max_norm\'], varargs=None, keywords=None, defaults=[\'None\', \'mean\', \'None\', \'None\', \'div\', \'None\'], "
argspec: "args=[\'embedding_weights\', \'sparse_ids\', \'sparse_weights\', \'combiner\', \'default_id\', \'name\', \'partition_strategy\', \'max_norm\', \'allow_fast_lookup\'], varargs=None, keywords=None, defaults=[\'None\', \'mean\', \'None\', \'None\', \'div\', \'None\', \'False\'], "
}
member_method {
name: "sampled_softmax_loss"
Expand Down
8 changes: 8 additions & 0 deletions tensorflow/tools/api/golden/v1/tensorflow.pbtxt
View file
Expand Up @@ -1848,6 +1848,14 @@ tf_module {
name: "quantized_concat"
argspec: "args=[\'concat_dim\', \'values\', \'input_mins\', \'input_maxes\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "ragged_fill_empty_rows"
argspec: "args=[\'value_rowids\', \'values\', \'nrows\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "ragged_fill_empty_rows_grad"
argspec: "args=[\'reverse_index_map\', \'grad_values\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "random_crop"
argspec: "args=[\'value\', \'size\', \'seed\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\'], "
Expand Down
8 changes: 8 additions & 0 deletions tensorflow/tools/api/golden/v1/tensorflow.raw_ops.pbtxt
View file
Expand Up @@ -3304,6 +3304,14 @@ tf_module {
name: "RaggedCross"
argspec: "args=[\'ragged_values\', \'ragged_row_splits\', \'sparse_indices\', \'sparse_values\', \'sparse_shape\', \'dense_inputs\', \'input_order\', \'hashed_output\', \'num_buckets\', \'hash_key\', \'out_values_type\', \'out_row_splits_type\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "RaggedFillEmptyRows"
argspec: "args=[\'value_rowids\', \'values\', \'nrows\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "RaggedFillEmptyRowsGrad"
argspec: "args=[\'reverse_index_map\', \'grad_values\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "RaggedGather"
argspec: "args=[\'params_nested_splits\', \'params_dense_values\', \'indices\', \'OUTPUT_RAGGED_RANK\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
Expand Down
4 changes: 2 additions & 2 deletions tensorflow/tools/api/golden/v2/tensorflow.nn.pbtxt
View file
Expand Up @@ -162,7 +162,7 @@ tf_module {
}
member_method {
name: "embedding_lookup_sparse"
argspec: "args=[\'params\', \'sp_ids\', \'sp_weights\', \'combiner\', \'max_norm\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'None\'], "
argspec: "args=[\'params\', \'sp_ids\', \'sp_weights\', \'combiner\', \'max_norm\', \'allow_fast_lookup\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'None\', \'False\', \'None\'], "
}
member_method {
name: "erosion2d"
Expand Down Expand Up @@ -270,7 +270,7 @@ tf_module {
}
member_method {
name: "safe_embedding_lookup_sparse"
argspec: "args=[\'embedding_weights\', \'sparse_ids\', \'sparse_weights\', \'combiner\', \'default_id\', \'max_norm\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'mean\', \'None\', \'None\', \'None\'], "
argspec: "args=[\'embedding_weights\', \'sparse_ids\', \'sparse_weights\', \'combiner\', \'default_id\', \'max_norm\', \'allow_fast_lookup\', \'name\'], varargs=None, keywords=None, defaults=[\'None\', \'mean\', \'None\', \'None\', \'False\', \'None\'], "
}
member_method {
name: "sampled_softmax_loss"
Expand Down
8 changes: 8 additions & 0 deletions tensorflow/tools/api/golden/v2/tensorflow.pbtxt
View file
Expand Up @@ -892,6 +892,14 @@ tf_module {
name: "py_function"
argspec: "args=[\'func\', \'inp\', \'Tout\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "ragged_fill_empty_rows"
argspec: "args=[\'value_rowids\', \'values\', \'nrows\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "ragged_fill_empty_rows_grad"
argspec: "args=[\'reverse_index_map\', \'grad_values\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "random_index_shuffle"
argspec: "args=[\'index\', \'seed\', \'max_index\', \'rounds\', \'name\'], varargs=None, keywords=None, defaults=[\'4\', \'None\'], "
Expand Down
8 changes: 8 additions & 0 deletions tensorflow/tools/api/golden/v2/tensorflow.raw_ops.pbtxt
View file
Expand Up @@ -3304,6 +3304,14 @@ tf_module {
name: "RaggedCross"
argspec: "args=[\'ragged_values\', \'ragged_row_splits\', \'sparse_indices\', \'sparse_values\', \'sparse_shape\', \'dense_inputs\', \'input_order\', \'hashed_output\', \'num_buckets\', \'hash_key\', \'out_values_type\', \'out_row_splits_type\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "RaggedFillEmptyRows"
argspec: "args=[\'value_rowids\', \'values\', \'nrows\', \'default_value\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "RaggedFillEmptyRowsGrad"
argspec: "args=[\'reverse_index_map\', \'grad_values\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
}
member_method {
name: "RaggedGather"
argspec: "args=[\'params_nested_splits\', \'params_dense_values\', \'indices\', \'OUTPUT_RAGGED_RANK\', \'name\'], varargs=None, keywords=None, defaults=[\'None\'], "
Expand Down

0 comments on commit 5165f07

Please sign in to comment.