Update init for sparse matrix op.

PiperOrigin-RevId: 471237779 Change-Id: Idd57748df2cc910ff80d79a6113ab68f982798fd
google-deepmind · Aug 31, 2022 · 51f5990 · 51f5990
1 parent b90a5aa
commit 51f5990
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/jraph/__init__.py b/jraph/__init__.py
@@ -64,6 +64,7 @@
 from jraph._src.utils import segment_softmax
 from jraph._src.utils import segment_sum
 from jraph._src.utils import segment_variance
+from jraph._src.utils import sparse_matrix_to_graphs_tuple
 from jraph._src.utils import unbatch
 from jraph._src.utils import unbatch_np
 from jraph._src.utils import unpad_with_graphs
@@ -91,7 +92,8 @@
            "segment_min_or_constant", "segment_softmax", "segment_sum",
            "partition_softmax", "concatenated_args",
            "get_fully_connected_graph", "dynamically_batch",
-           "with_zero_out_padding_outputs", "zero_out_padding")
+           "with_zero_out_padding_outputs", "zero_out_padding",
+           "sparse_matrix_to_graphs_tuple")
 
 #  _________________________________________
 # / Please don't use symbols in `_src` they \