tensorflow - tf.sparse.SparseTensor 上的 Argmax

Question

给定一个形状为 [B, d, d', v] 的稀疏张量 x，我如何tf.argmax(x, axis=2)在保持稀疏性的同时进行计算？

x = tf.sparse.SparseTensor(values=..., indices=..., shape=[B, d1, d2, v])
# How to implement something like
sparse_argmax(x, axis=2)

提前致谢！

标签： tensorflowtensorflow2.0sparse-matrix