Adds typing information to the nn.loss package.

PiperOrigin-RevId: 420046706
tensorflow · Jan 6, 2022 · 3b030f8 · 3b030f8
1 parent bab2352
commit 3b030f8
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 2 deletions.
diff --git a/tensorflow_graphics/nn/loss/chamfer_distance.py b/tensorflow_graphics/nn/loss/chamfer_distance.py
@@ -21,9 +21,12 @@
 
 from tensorflow_graphics.util import export_api
 from tensorflow_graphics.util import shape
+from tensorflow_graphics.util import type_alias
 
 
-def evaluate(point_set_a, point_set_b, name="chamfer_distance_evaluate"):
+def evaluate(point_set_a: type_alias.TensorLike,
+             point_set_b: type_alias.TensorLike,
+             name: str = "chamfer_distance_evaluate") -> tf.Tensor:
   """Computes the Chamfer distance for the given two point sets.
 
   Note:

diff --git a/tensorflow_graphics/nn/loss/hausdorff_distance.py b/tensorflow_graphics/nn/loss/hausdorff_distance.py
@@ -21,9 +21,12 @@
 
 from tensorflow_graphics.util import export_api
 from tensorflow_graphics.util import shape
+from tensorflow_graphics.util import type_alias
 
 
-def evaluate(point_set_a, point_set_b, name="hausdorff_distance_evaluate"):
+def evaluate(point_set_a: type_alias.TensorLike,
+             point_set_b: type_alias.TensorLike,
+             name: str = "hausdorff_distance_evaluate") -> tf.Tensor:
   """Computes the Hausdorff distance from point_set_a to point_set_b.
 
   Note: