Fix error in `sparse_federated_learning.ipynb`.
The correct dense shape after global token to local token mapping should have MAX_TOKENS_SELECTED_PER_CLIENT in the second value.
" st = tf.sparse.SparseTensor(\n",
" updated_indices,\n",
" tf.ones(tf.size(tokens), dtype=FEATURE_DTYPE),\n",
- " dense_shape=sparse_tokens.dense_shape)\n",
+ " # Each client has at most MAX_TOKENS_SELECTED_PER_CLIENT distinct tokens.\n",
+ " dense_shape=[sparse_tokens.dense_shape[0], MAX_TOKENS_SELECTED_PER_CLIENT])\n",
" st = tf.sparse.reorder(st)\n",
" return st\n",
"\n",
