Unify the internal and external tutorials.

dgf/src/io/dataset_loader_traffic.py

@@ -137,8 +137,8 @@ def traffic_cns_name(dataset: str, forecast_horizon_seconds: int) -> str:
   if forecast_horizon_seconds not in CACHED_FORECAST_HORIZONS_SECONDS:
     raise ValueError(
         f"Forecast horizon {forecast_horizon_seconds}s is not cached on CNS for"
-        f" dataset {dataset!r}. Cached horizons are: "
-        f"{sorted(CACHED_FORECAST_HORIZONS_SECONDS)} (seconds, corresponding"
+        f" dataset {dataset!r}. Cached horizons are:"
+        f" {sorted(CACHED_FORECAST_HORIZONS_SECONDS)} (seconds, corresponding"
         f" to {[CACHED_FORECAST_HORIZONS_SECONDS[h] for h in sorted(CACHED_FORECAST_HORIZONS_SECONDS)]})."
         " To build a graph with an arbitrary horizon, pass"
         " repo=dgf.io.Repo.WEB."
@@ -278,17 +278,13 @@ def download_traffic_sensor_graph(
       ["sensor_id", "latitude", "longitude"],
       spec.locations_have_header,
   )
-  distances = distances.assign(
-      **{
-          "from": [_decode_sensor_id(value) for value in distances["from"]],
-          "to": [_decode_sensor_id(value) for value in distances["to"]],
-          "cost": distances["cost"].astype(np.float32),
-      }
-  )
+  distances = distances.assign(**{
+      "from": [_decode_sensor_id(value) for value in distances["from"]],
+      "to": [_decode_sensor_id(value) for value in distances["to"]],
+      "cost": distances["cost"].astype(np.float32),
+  })
   locations = locations.assign(
-      sensor_id=[
-          _decode_sensor_id(value) for value in locations["sensor_id"]
-      ],
+      sensor_id=[_decode_sensor_id(value) for value in locations["sensor_id"]],
       latitude=locations["latitude"].astype(np.float32),
       longitude=locations["longitude"].astype(np.float32),
   )
@@ -465,12 +461,12 @@ def build_traffic_graph(
       ),
       "time": sensor_times,
       "speed": sensor_speeds,
-      "latitude": locations["latitude"].to_numpy(dtype=np.float32)[
-          location_positions
-      ],
-      "longitude": locations["longitude"].to_numpy(dtype=np.float32)[
-          location_positions
-      ],
+      "latitude": (
+          locations["latitude"].to_numpy(dtype=np.float32)[location_positions]
+      ),
+      "longitude": (
+          locations["longitude"].to_numpy(dtype=np.float32)[location_positions]
+      ),
   }
   sensor_schema_features: dict[str, schema_lib.FeatureSchema] = {
       "#id": schema_lib.FeatureSchema(
@@ -541,8 +537,9 @@ def build_traffic_graph(
     )
 
   query_times = np.repeat(timestamps[query_time_indices], num_sensors)
-  query_sensors = np.tile(np.arange(num_sensors, dtype=np.int64),
-                          num_query_times)
+  query_sensors = np.tile(
+      np.arange(num_sensors, dtype=np.int64), num_query_times
+  )
   target_speeds = values[target_time_indices].reshape(-1)
 
   # Chronological 70/10/20 split, matching the canonical METR-LA and PEMS-BAY
@@ -677,7 +674,7 @@ def fetch_traffic_graph(
     adjacency_threshold: float = 0.1,
     repo: Repo | str = Repo.AUTO,
 ) -> tuple[in_memory_graph_lib.InMemoryGraph, schema_lib.GraphSchema]:
-  """Downloads and loads a traffic speed forecasting benchmark into memory.
+  """Gets the METR-LA and PEMS-BAY traffic speed forecasting datasets.
 
   Both supported datasets record the speed of highway loop detectors every five
   minutes: METR-LA covers 207 detectors of the Los Angeles county highways over

dgf/src/transform/merge.py

@@ -50,6 +50,8 @@ class GraphMerger:
   (e.g., the padding has space for 100 nodes but 120 nodes are provided),
   raises an `InsufficientPaddingError` exception.
 
+  If `padding` is provided, sentinel nodes will addeds.
+
   Attributes:
     schema: The original input GraphSchema.
     padding: The padding configuration to apply.

doc/docs/api/dgf-analyse.md

@@ -11,9 +11,21 @@ hide:
 ## dgf.analyse.feature_statistics_from_graphs  # {: #section-feature-statistics-from-graphs}
 ::: dgf.analyse.feature_statistics_from_graphs
 
+## dgf.analyse.infer_schema_semantic  # {: #section-infer-schema-semantic}
+::: dgf.analyse.infer_schema_semantic
+
+## dgf.analyse.make_histogram  # {: #section-make-histogram}
+::: dgf.analyse.make_histogram
+
 ## dgf.analyse.padding_from_graph_generator  # {: #section-padding-from-graph-generator}
 ::: dgf.analyse.padding_from_graph_generator
 
 ## dgf.analyse.print_schema  # {: #section-print-schema}
 ::: dgf.analyse.print_schema
 
+## dgf.analyse.topology_statistics  # {: #section-topology-statistics}
+::: dgf.analyse.topology_statistics
+
+## dgf.analyse.topology_statistics_from_graphs  # {: #section-topology-statistics-from-graphs}
+::: dgf.analyse.topology_statistics_from_graphs
+

doc/docs/api/dgf-beam-io.md

@@ -44,3 +44,6 @@ hide:
 ## dgf.beam.io.write_tfgnn_graphs  # {: #section-write-tfgnn-graphs}
 ::: dgf.beam.io.write_tfgnn_graphs
 
+## dgf.beam.io.write_topology_statistics  # {: #section-write-topology-statistics}
+::: dgf.beam.io.write_topology_statistics
+

doc/docs/api/dgf-convert.md

@@ -14,6 +14,9 @@ hide:
 ## dgf.convert.graph_to_networkx  # {: #section-graph-to-networkx}
 ::: dgf.convert.graph_to_networkx
 
+## dgf.convert.graph_to_pyg_data  # {: #section-graph-to-pyg-data}
+::: dgf.convert.graph_to_pyg_data
+
 ## dgf.convert.graph_to_serialized_tfgnn_graph  # {: #section-graph-to-serialized-tfgnn-graph}
 ::: dgf.convert.graph_to_serialized_tfgnn_graph
 
@@ -41,9 +44,15 @@ hide:
 ## dgf.convert.schema_to_sparse_deferred_schema  # {: #section-schema-to-sparse-deferred-schema}
 ::: dgf.convert.schema_to_sparse_deferred_schema
 
+## dgf.convert.schema_to_tfgnn_graph_parsing_spec  # {: #section-schema-to-tfgnn-graph-parsing-spec}
+::: dgf.convert.schema_to_tfgnn_graph_parsing_spec
+
 ## dgf.convert.schema_to_tfgnn_schema  # {: #section-schema-to-tfgnn-schema}
 ::: dgf.convert.schema_to_tfgnn_schema
 
+## dgf.convert.serialized_tfgnn_graph_to_tf_graph  # {: #section-serialized-tfgnn-graph-to-tf-graph}
+::: dgf.convert.serialized_tfgnn_graph_to_tf_graph
+
 ## dgf.convert.sparse_deferred_struct_to_graph  # {: #section-sparse-deferred-struct-to-graph}
 ::: dgf.convert.sparse_deferred_struct_to_graph
 
@@ -53,6 +62,9 @@ hide:
 ## dgf.convert.tf_graph_to_tf_graph_dict  # {: #section-tf-graph-to-tf-graph-dict}
 ::: dgf.convert.tf_graph_to_tf_graph_dict
 
+## dgf.convert.tfgnn_graph_dict_to_tf_graph  # {: #section-tfgnn-graph-dict-to-tf-graph}
+::: dgf.convert.tfgnn_graph_dict_to_tf_graph
+
 ## dgf.convert.tfgnn_graph_to_graph  # {: #section-tfgnn-graph-to-graph}
 ::: dgf.convert.tfgnn_graph_to_graph
 

doc/docs/api/dgf-data.md

@@ -14,6 +14,9 @@ hide:
 ## dgf.data.FeatureFormat  # {: #section-featureformat}
 ::: dgf.data.FeatureFormat
 
+## dgf.data.FeaturePadding  # {: #section-featurepadding}
+::: dgf.data.FeaturePadding
+
 ## dgf.data.FeatureSchema  # {: #section-featureschema}
 ::: dgf.data.FeatureSchema
 
@@ -38,6 +41,18 @@ hide:
 ## dgf.data.GraphSchemaV2  # {: #section-graphschemav2}
 ::: dgf.data.GraphSchemaV2
 
+## dgf.data.GraphSnapshots  # {: #section-graphsnapshots}
+::: dgf.data.GraphSnapshots
+
+## dgf.data.GraphSnapshotsFormat  # {: #section-graphsnapshotsformat}
+::: dgf.data.GraphSnapshotsFormat
+
+## dgf.data.GraphSnapshotsMetadata  # {: #section-graphsnapshotsmetadata}
+::: dgf.data.GraphSnapshotsMetadata
+
+## dgf.data.Histogram  # {: #section-histogram}
+::: dgf.data.Histogram
+
 ## dgf.data.InMemoryEdgeSet  # {: #section-inmemoryedgeset}
 ::: dgf.data.InMemoryEdgeSet
 

doc/docs/api/dgf-filesystem.md

@@ -23,6 +23,9 @@ hide:
 ## dgf.filesystem.open_read  # {: #section-open-read}
 ::: dgf.filesystem.open_read
 
+## dgf.filesystem.open_write  # {: #section-open-write}
+::: dgf.filesystem.open_write
+
 ## dgf.filesystem.remove_paths  # {: #section-remove-paths}
 ::: dgf.filesystem.remove_paths
 

doc/docs/api/dgf-io.md

@@ -17,9 +17,15 @@ hide:
 ## dgf.io.fetch_graphland_graph  # {: #section-fetch-graphland-graph}
 ::: dgf.io.fetch_graphland_graph
 
+## dgf.io.fetch_jena_climate_graph  # {: #section-fetch-jena-climate-graph}
+::: dgf.io.fetch_jena_climate_graph
+
 ## dgf.io.fetch_ogb_graph  # {: #section-fetch-ogb-graph}
 ::: dgf.io.fetch_ogb_graph
 
+## dgf.io.fetch_traffic_graph  # {: #section-fetch-traffic-graph}
+::: dgf.io.fetch_traffic_graph
+
 ## dgf.io.read_bigquery_graph  # {: #section-read-bigquery-graph}
 ::: dgf.io.read_bigquery_graph
 
@@ -32,12 +38,21 @@ hide:
 ## dgf.io.read_graph  # {: #section-read-graph}
 ::: dgf.io.read_graph
 
+## dgf.io.read_graph_snapshots  # {: #section-read-graph-snapshots}
+::: dgf.io.read_graph_snapshots
+
+## dgf.io.read_graph_snapshots_as_temporal_graph  # {: #section-read-graph-snapshots-as-temporal-graph}
+::: dgf.io.read_graph_snapshots_as_temporal_graph
+
 ## dgf.io.read_graphai_hgraph  # {: #section-read-graphai-hgraph}
 ::: dgf.io.read_graphai_hgraph
 
 ## dgf.io.read_schema  # {: #section-read-schema}
 ::: dgf.io.read_schema
 
+## dgf.io.read_snapshot_metadata  # {: #section-read-snapshot-metadata}
+::: dgf.io.read_snapshot_metadata
+
 ## dgf.io.read_spanner_graph  # {: #section-read-spanner-graph}
 ::: dgf.io.read_spanner_graph
 
@@ -50,6 +65,9 @@ hide:
 ## dgf.io.read_tfgnn_graphs  # {: #section-read-tfgnn-graphs}
 ::: dgf.io.read_tfgnn_graphs
 
+## dgf.io.read_topology_statistics  # {: #section-read-topology-statistics}
+::: dgf.io.read_topology_statistics
+
 ## dgf.io.write_feature_statistics  # {: #section-write-feature-statistics}
 ::: dgf.io.write_feature_statistics
 
@@ -59,9 +77,15 @@ hide:
 ## dgf.io.write_schema  # {: #section-write-schema}
 ::: dgf.io.write_schema
 
+## dgf.io.write_snapshot_metadata  # {: #section-write-snapshot-metadata}
+::: dgf.io.write_snapshot_metadata
+
 ## dgf.io.write_text_proto  # {: #section-write-text-proto}
 ::: dgf.io.write_text_proto
 
 ## dgf.io.write_tfgnn_graphs  # {: #section-write-tfgnn-graphs}
 ::: dgf.io.write_tfgnn_graphs
 
+## dgf.io.write_topology_statistics  # {: #section-write-topology-statistics}
+::: dgf.io.write_topology_statistics
+

doc/docs/api/dgf-jax-layers.md

@@ -20,6 +20,12 @@ hide:
 ## dgf.jax.layers.EmbedAndHomogenizeGraphConfig  # {: #section-embedandhomogenizegraphconfig}
 ::: dgf.jax.layers.EmbedAndHomogenizeGraphConfig
 
+## dgf.jax.layers.EmbedFeatureGroups  # {: #section-embedfeaturegroups}
+::: dgf.jax.layers.EmbedFeatureGroups
+
+## dgf.jax.layers.EmbedFeatureGroupsConfig  # {: #section-embedfeaturegroupsconfig}
+::: dgf.jax.layers.EmbedFeatureGroupsConfig
+
 ## dgf.jax.layers.EmbedFeatureSet  # {: #section-embedfeatureset}
 ::: dgf.jax.layers.EmbedFeatureSet
 
@@ -50,6 +56,9 @@ hide:
 ## dgf.jax.layers.GenericBlockConfig  # {: #section-genericblockconfig}
 ::: dgf.jax.layers.GenericBlockConfig
 
+## dgf.jax.layers.GnnPlus  # {: #section-gnnplus}
+::: dgf.jax.layers.GnnPlus
+
 ## dgf.jax.layers.HeterogeneousGraphAttentionNetwork  # {: #section-heterogeneousgraphattentionnetwork}
 ::: dgf.jax.layers.HeterogeneousGraphAttentionNetwork
 

doc/docs/api/dgf-jax.md

@@ -5,18 +5,6 @@ hide:
 
 # dgf.jax
 
-## dgf.jax.JaxBaseConfig  # {: #section-jaxbaseconfig}
-::: dgf.jax.JaxBaseConfig
-
-## dgf.jax.get_activation  # {: #section-get-activation}
-::: dgf.jax.get_activation
-
-## dgf.jax.jnp_dtype_from_string  # {: #section-jnp-dtype-from-string}
-::: dgf.jax.jnp_dtype_from_string
-
-## dgf.jax.jnp_name_from_dtype  # {: #section-jnp-name-from-dtype}
-::: dgf.jax.jnp_name_from_dtype
-
 ## dgf.jax.train  # {: #section-train}
 ::: dgf.jax.train
 

doc/docs/api/dgf-learning.md

@@ -14,6 +14,9 @@ hide:
 ## dgf.learning.NodePredictionModel  # {: #section-nodepredictionmodel}
 ::: dgf.learning.NodePredictionModel
 
+## dgf.learning.TFFunctionInputFormat  # {: #section-tffunctioninputformat}
+::: dgf.learning.TFFunctionInputFormat
+
 ## dgf.learning.load_model  # {: #section-load-model}
 ::: dgf.learning.load_model
 

doc/docs/api/dgf-sampling.md

@@ -26,6 +26,9 @@ hide:
 ## dgf.sampling.extract_beam_nodes_ids  # {: #section-extract-beam-nodes-ids}
 ::: dgf.sampling.extract_beam_nodes_ids
 
+## dgf.sampling.offline_distributed_sampler_gcp  # {: #section-offline-distributed-sampler-gcp}
+::: dgf.sampling.offline_distributed_sampler_gcp
+
 ## dgf.sampling.sample_with_beam_semi_distributed_sampler  # {: #section-sample-with-beam-semi-distributed-sampler}
 ::: dgf.sampling.sample_with_beam_semi_distributed_sampler
 

doc/docs/api/dgf-transform.md

@@ -8,6 +8,12 @@ hide:
 ## dgf.transform.AutoNormalizeConfig  # {: #section-autonormalizeconfig}
 ::: dgf.transform.AutoNormalizeConfig
 
+## dgf.transform.CalendarFeature  # {: #section-calendarfeature}
+::: dgf.transform.CalendarFeature
+
+## dgf.transform.CalendarNormalizer  # {: #section-calendarnormalizer}
+::: dgf.transform.CalendarNormalizer
+
 ## dgf.transform.ContainsLabelPredicate  # {: #section-containslabelpredicate}
 ::: dgf.transform.ContainsLabelPredicate
 
@@ -17,6 +23,9 @@ hide:
 ## dgf.transform.GNNDatasetPreparator  # {: #section-gnndatasetpreparator}
 ::: dgf.transform.GNNDatasetPreparator
 
+## dgf.transform.GraphMerger  # {: #section-graphmerger}
+::: dgf.transform.GraphMerger
+
 ## dgf.transform.GraphNormalizer  # {: #section-graphnormalizer}
 ::: dgf.transform.GraphNormalizer
 
@@ -29,12 +38,18 @@ hide:
 ## dgf.transform.NumNodesPredicate  # {: #section-numnodespredicate}
 ::: dgf.transform.NumNodesPredicate
 
+## dgf.transform.SequentialNormalizer  # {: #section-sequentialnormalizer}
+::: dgf.transform.SequentialNormalizer
+
 ## dgf.transform.SinusoidTimedeltaNormalizer  # {: #section-sinusoidtimedeltanormalizer}
 ::: dgf.transform.SinusoidTimedeltaNormalizer
 
 ## dgf.transform.SoftQuantileNormalizer  # {: #section-softquantilenormalizer}
 ::: dgf.transform.SoftQuantileNormalizer
 
+## dgf.transform.TimedeltaNormalizer  # {: #section-timedeltanormalizer}
+::: dgf.transform.TimedeltaNormalizer
+
 ## dgf.transform.apply_feature  # {: #section-apply-feature}
 ::: dgf.transform.apply_feature
 
@@ -44,6 +59,9 @@ hide:
 ## dgf.transform.batch_indices_generator  # {: #section-batch-indices-generator}
 ::: dgf.transform.batch_indices_generator
 
+## dgf.transform.combine_graph_snapshots  # {: #section-combine-graph-snapshots}
+::: dgf.transform.combine_graph_snapshots
+
 ## dgf.transform.drop_edge_features  # {: #section-drop-edge-features}
 ::: dgf.transform.drop_edge_features
 
@@ -65,9 +83,6 @@ hide:
 ## dgf.transform.homogenize  # {: #section-homogenize}
 ::: dgf.transform.homogenize
 
-## dgf.transform.merge_graphs  # {: #section-merge-graphs}
-::: dgf.transform.merge_graphs
-
 ## dgf.transform.propagate_timestamp_to_edges  # {: #section-propagate-timestamp-to-edges}
 ::: dgf.transform.propagate_timestamp_to_edges
 

doc/docs/api/dgf-validate.md

@@ -5,6 +5,12 @@ hide:
 
 # dgf.validate
 
+## dgf.validate.fix_schema  # {: #section-fix-schema}
+::: dgf.validate.fix_schema
+
 ## dgf.validate.validate_graph  # {: #section-validate-graph}
 ::: dgf.validate.validate_graph
 
+## dgf.validate.validate_snapshots  # {: #section-validate-snapshots}
+::: dgf.validate.validate_snapshots
+