Adds typing information to the package math.

PiperOrigin-RevId: 414426546

Author: G4G

tensorflow_graphics/math/feature_representation.py

@@ -20,11 +20,12 @@
 import tensorflow as tf
 
 from tensorflow_graphics.util import export_api
+from tensorflow_graphics.util.type_alias import TensorLike
 
 
-def positional_encoding(features: tf.Tensor,
+def positional_encoding(features: TensorLike,
                         num_frequencies: int,
-                        name="positional_encoding") -> tf.Tensor:
+                        name: str = "positional_encoding") -> TensorLike:
   """Positional enconding of a tensor as described in the NeRF paper (https://arxiv.org/abs/2003.08934).
 
   Args:

tensorflow_graphics/math/math_helpers.py

@@ -19,17 +19,18 @@
 
 import numpy as np
 import tensorflow as tf
-
 from tensorflow_graphics.util import asserts
 from tensorflow_graphics.util import export_api
 from tensorflow_graphics.util import safe_ops
 from tensorflow_graphics.util import shape
+from tensorflow_graphics.util.type_alias import Float
+from tensorflow_graphics.util.type_alias import TensorLike
 
 
-def cartesian_to_spherical_coordinates(point_cartesian,
-                                       eps=None,
-                                       name="cartesian_to_spherical_coordinates"
-                                      ):
+def cartesian_to_spherical_coordinates(
+    point_cartesian: TensorLike,
+    eps: Float = None,
+    name: str = "cartesian_to_spherical_coordinates") -> tf.Tensor:
   """Function to transform Cartesian coordinates to spherical coordinates.
 
   This function assumes a right handed coordinate system with `z` pointing up.
@@ -77,7 +78,7 @@ def _double_factorial_loop_condition(n, result, two):
   return tf.cast(tf.math.count_nonzero(tf.greater_equal(n, two)), tf.bool)
 
 
-def double_factorial(n):
+def double_factorial(n: TensorLike) -> TensorLike:
   """Computes the double factorial of `n`.
 
   Note:
@@ -100,7 +101,7 @@ def double_factorial(n):
   return result
 
 
-def factorial(n):
+def factorial(n: TensorLike) -> TensorLike:
   """Computes the factorial of `n`.
 
   Note:
@@ -117,9 +118,9 @@ def factorial(n):
   return tf.exp(tf.math.lgamma(n + 1))
 
 
-def spherical_to_cartesian_coordinates(point_spherical,
-                                       name="spherical_to_cartesian_coordinates"
-                                      ):
+def spherical_to_cartesian_coordinates(
+    point_spherical: TensorLike,
+    name: str = "spherical_to_cartesian_coordinates") -> TensorLike:
   """Function to transform Cartesian coordinates to spherical coordinates.
 
   Note:
@@ -156,9 +157,9 @@ def spherical_to_cartesian_coordinates(point_spherical,
     return tf.stack((x, y, z), axis=-1)
 
 
-def square_to_spherical_coordinates(point_2d,
-                                    name="math_square_to_spherical_coordinates"
-                                   ):
+def square_to_spherical_coordinates(
+    point_2d: TensorLike,
+    name: str = "math_square_to_spherical_coordinates") -> TensorLike:
   """Maps points from a unit square to a unit sphere.
 
   Note:

tensorflow_graphics/math/spherical_harmonics.py

@@ -17,6 +17,8 @@
 from __future__ import division
 from __future__ import print_function
 
+from typing import Tuple
+
 import numpy as np
 from six.moves import range
 import tensorflow as tf
@@ -26,12 +28,14 @@
 from tensorflow_graphics.util import asserts
 from tensorflow_graphics.util import export_api
 from tensorflow_graphics.util import shape
+from tensorflow_graphics.util.type_alias import TensorLike
 
 
-def integration_product(harmonics1,
-                        harmonics2,
-                        keepdims=True,
-                        name="spherical_harmonics_convolution"):
+def integration_product(
+    harmonics1: TensorLike,
+    harmonics2: TensorLike,
+    keepdims: bool = True,
+    name: str = "spherical_harmonics_convolution") -> TensorLike:
   """Computes the integral of harmonics1.harmonics2 over the sphere.
 
   Note:
@@ -72,7 +76,8 @@ def integration_product(harmonics1,
 
 
 def generate_l_m_permutations(
-    max_band, name="spherical_harmonics_generate_l_m_permutations"):
+    max_band: int,
+    name: str = "spherical_harmonics_generate_l_m_permutations") -> Tuple[TensorLike, TensorLike]:  # pylint: disable=line-too-long
   """Generates permutations of degree l and order m for spherical harmonics.
 
   Args:
@@ -94,7 +99,9 @@ def generate_l_m_permutations(
             tf.convert_to_tensor(value=order_m))
 
 
-def generate_l_m_zonal(max_band, name="spherical_harmonics_generate_l_m_zonal"):
+def generate_l_m_zonal(
+    max_band: int,
+    name: str = "spherical_harmonics_generate_l_m_zonal") -> Tuple[TensorLike, TensorLike]:  # pylint: disable=line-too-long
   """Generates l and m coefficients for zonal harmonics.
 
   Args:
@@ -154,7 +161,9 @@ def _evaluate_legendre_polynomial_branch(l, m, x, pmm):
   return res
 
 
-def evaluate_legendre_polynomial(degree_l, order_m, x):
+def evaluate_legendre_polynomial(degree_l: TensorLike,
+                                 order_m: TensorLike,
+                                 x: TensorLike) -> TensorLike:
   """Evaluates the Legendre polynomial of degree l and order m at x.
 
   Note:
@@ -227,11 +236,11 @@ def _evaluate_spherical_harmonics_branch(degree,
 
 
 def evaluate_spherical_harmonics(
-    degree_l,
-    order_m,
-    theta,
-    phi,
-    name="spherical_harmonics_evaluate_spherical_harmonics"):
+    degree_l: TensorLike,
+    order_m: TensorLike,
+    theta: TensorLike,
+    phi: TensorLike,
+    name: str = "spherical_harmonics_evaluate_spherical_harmonics") -> TensorLike:    # pylint: disable=line-too-long
   """Evaluates a point sample of a Spherical Harmonic basis function.
 
   Note:
@@ -305,10 +314,11 @@ def evaluate_spherical_harmonics(
     return tf.where(tf.equal(order_m, zeros), result_m_zero, result_branch)
 
 
-def rotate_zonal_harmonics(zonal_coeffs,
-                           theta,
-                           phi,
-                           name="spherical_harmonics_rotate_zonal_harmonics"):
+def rotate_zonal_harmonics(
+    zonal_coeffs: TensorLike,
+    theta: TensorLike,
+    phi: TensorLike,
+    name: str = "spherical_harmonics_rotate_zonal_harmonics") -> TensorLike:
   """Rotates zonal harmonics.
 
   Note:
@@ -356,8 +366,9 @@ def rotate_zonal_harmonics(zonal_coeffs,
         l_broadcasted, m_broadcasted, theta, phi)
 
 
-def tile_zonal_coefficients(coefficients,
-                            name="spherical_harmonics_tile_zonal_coefficients"):
+def tile_zonal_coefficients(
+    coefficients: TensorLike,
+    name: str = "spherical_harmonics_tile_zonal_coefficients") -> TensorLike:
   """Tiles zonal coefficients.
 
   Zonal Harmonics only contains the harmonics where m=0. This function returns

tensorflow_graphics/math/vector.py

@@ -22,9 +22,13 @@
 from tensorflow_graphics.util import asserts
 from tensorflow_graphics.util import export_api
 from tensorflow_graphics.util import shape
+from tensorflow_graphics.util.type_alias import TensorLike
 
 
-def cross(vector1, vector2, axis=-1, name="vector_cross"):
+def cross(vector1: TensorLike,
+          vector2: TensorLike,
+          axis: int = -1,
+          name: str = "vector_cross") -> TensorLike:
   """Computes the cross product between two tensors along an axis.
 
   Note:
@@ -62,7 +66,11 @@ def cross(vector1, vector2, axis=-1, name="vector_cross"):
     return tf.stack((n_x, n_y, n_z), axis=axis)
 
 
-def dot(vector1, vector2, axis=-1, keepdims=True, name="vector_dot"):
+def dot(vector1: TensorLike,
+        vector2: TensorLike,
+        axis: int = -1,
+        keepdims: bool = True,
+        name: str = "vector_dot") -> TensorLike:
   """Computes the dot product between two tensors along an axis.
 
   Note:
@@ -97,7 +105,10 @@ def dot(vector1, vector2, axis=-1, keepdims=True, name="vector_dot"):
         input_tensor=vector1 * vector2, axis=axis, keepdims=keepdims)
 
 
-def reflect(vector, normal, axis=-1, name="vector_reflect"):
+def reflect(vector: TensorLike,
+            normal: TensorLike,
+            axis: int = -1,
+            name: str = "vector_reflect") -> TensorLike:
   r"""Computes the reflection direction for an incident vector.
 
   For an incident vector \\(\mathbf{v}\\) and normal $$\mathbf{n}$$ this