Add dynamic API for math functions + tests. (#6066)

* Add math functions and tests.
* Move arithm op arguments to GPU if there's any GPU argument.
* Forbid mixing of DALI CPU and GPU tensors/batches in arithmetic ops and math functions.

---------
Signed-off-by: Michal Zientkiewicz <[email protected]>

Author: mzient

dali/python/nvidia/dali/experimental/dynamic/__init__.py

@@ -25,6 +25,7 @@
 
 from . import _fn
 from . import ops
+from . import math  # noqa: F401
 
 ops._initialize()
 _fn._initialize()

dali/python/nvidia/dali/experimental/dynamic/_batch.py

@@ -101,7 +101,30 @@ def _arithm_op(name, *args, **kwargs):
     from . import arithmetic_generic_op
 
     argsstr = " ".join(f"&{i}" for i in range(len(args)))
-    return arithmetic_generic_op(*args, expression_desc=f"{name}({argsstr})")
+    gpu = False
+    new_args = [None] * len(args)
+    for i, a in enumerate(args):
+        if isinstance(a, (Batch, Tensor)):
+            if a.device.device_type == "gpu":
+                gpu = True
+        else:
+            # TODO(michalz): We might use some caching here for common values.
+            if new_args is None:
+                new_args = list(args)
+            if gpu:
+                new_args[i] = _as_tensor(a, device="gpu")
+            else:
+                new_args[i] = _as_tensor(a)
+                if new_args[i].device.device_type == "gpu":
+                    gpu = True
+
+    for i in range(len(args)):
+        if new_args[i] is None:
+            if (args[i].device.device_type == "gpu") != gpu:
+                raise ValueError("Cannot mix GPU and CPU inputs.")
+            new_args[i] = args[i]
+
+    return arithmetic_generic_op(*new_args, expression_desc=f"{name}({argsstr})")
 
 
 class _TensorList:

dali/python/nvidia/dali/experimental/dynamic/math.py

@@ -0,0 +1,272 @@
+# Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from ._batch import _arithm_op
+
+
+def sqrt(input):
+    """Computes square root of values in `input`.
+
+    :rtype: Tensor or Batch of sqrt(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("sqrt", input)
+
+
+def rsqrt(input):
+    """Computes reciprocal of the square root of values in `input`.
+
+    :rtype: Tensor or Batch of rsqrt(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("rsqrt", input)
+
+
+def cbrt(input):
+    """Computes cube root of values in `input`.
+
+    :rtype: Tensor or Batch of cbrt(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("cbrt", input)
+
+
+def exp(input):
+    """Computes exponential of values in `input`.
+
+    :rtype: Tensor or Batch of exp(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("exp", input)
+
+
+def log(input):
+    """Computes natural logarithm (base-e) of values in `input`.
+
+    :rtype: Tensor or Batch of log(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("log", input)
+
+
+def log2(input):
+    """Computes logarithm (base-2) of values in `input`.
+
+    :rtype: Tensor or Batch of log2(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("log2", input)
+
+
+def log10(input):
+    """Computes logarithm (base-10) of values in `input`.
+
+    :rtype: Tensor or Batch of log10(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("log10", input)
+
+
+def abs(input):
+    """Computes absolute value of values in `input`.
+
+    :rtype: Tensor or Batch of abs(input). The type is preserved.
+    """
+    return _arithm_op("abs", input)
+
+
+def fabs(input):
+    """Computes float absolute value of values in `input`.
+
+    :rtype: Tensor or Batch of fabs(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("fabs", input)
+
+
+def floor(input):
+    """Computes floor of values in `input`.
+
+    :rtype: Tensor or Batch of floor(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("floor", input)
+
+
+def ceil(input):
+    """Computes ceil of values in `input`.
+
+    :rtype: Tensor or Batch of ceil(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("ceil", input)
+
+
+def sin(input):
+    """Computes sine of values in `input`.
+
+    :rtype: Tensor or Batch of sin(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("sin", input)
+
+
+def cos(input):
+    """Computes cosine of values in `input`.
+
+    :rtype: Tensor or Batch of cos(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("cos", input)
+
+
+def tan(input):
+    """Computes tangent of values in `input`.
+
+    :rtype: Tensor or Batch of tan(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("tan", input)
+
+
+def asin(input):
+    """Computes arcus sine of values in `input`.
+
+    :rtype: Tensor or Batch of asin(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("asin", input)
+
+
+def acos(input):
+    """Computes arcus cosine of values in `input`.
+
+    :rtype: Tensor or Batch of acos(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("acos", input)
+
+
+def atan(input):
+    """Computes arcus tangent of values in `input`.
+
+    :rtype: Tensor or Batch of atan(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("atan", input)
+
+
+def sinh(input):
+    """Computes hyperbolic sine of values in `input`.
+
+    :rtype: Tensor or Batch of sinh(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("sinh", input)
+
+
+def cosh(input):
+    """Computes hyperbolic cosine of values in `input`.
+
+    :rtype: Tensor or Batch of cosh(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("cosh", input)
+
+
+def tanh(input):
+    """Computes hyperbolic tangent of values in `input`.
+
+    :rtype: Tensor or Batch of tanh(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("tanh", input)
+
+
+def asinh(input):
+    """Computes inverse hyperbolic sine of values in `input`.
+
+    :rtype: Tensor or Batch of asinh(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("asinh", input)
+
+
+def acosh(input):
+    """Computes inverse hyperbolic cosine of values in `input`.
+
+    :rtype: Tensor or Batch of acosh(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("acosh", input)
+
+
+def atanh(input):
+    """Computes inverse hyperbolic tangent of values in `input`.
+
+    :rtype: Tensor or Batch of atanh(input). If input is an integer, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("atanh", input)
+
+
+def min(left, right):
+    """Computes minima of corresponding values in `left` and `right`.
+
+    :rtype: Tensor or Batch of the type that is calculated based on the type promotion rules.
+    """
+    return _arithm_op("min", left, right)
+
+
+def max(left, right):
+    """Computes maxima of corresponding values in `left` and `right`.
+
+    :rtype: Tensor or Batch of the type that is calculated based on the type promotion rules.
+    """
+    return _arithm_op("max", left, right)
+
+
+def pow(base, exponent):
+    """Computes base to the power of exponents, that is base ** exponent.
+
+    :rtype: Tensor or Batch of pow(base, exponent). Type is calculated based on the type
+            promotion rules.
+    """
+    return _arithm_op("pow", base, exponent)
+
+
+def fpow(base, exponent):
+    """Computes base to the power of exponents as floating point numbers.
+
+    :rtype: Tensor or Batch of pow(base, exponent). If all inputs are integers, the result
+            will be float, otherwise the type is preserved.
+    """
+    return _arithm_op("fpow", base, exponent)
+
+
+def atan2(x, y):
+    """Computes arcus tangent of corresponding values in  x / y.
+
+    :rtype: Tensor or Batch of atan2(x, y). If all inputs are integers, the result will be float,
+            otherwise the type is preserved.
+    """
+    return _arithm_op("atan2", x, y)
+
+
+def clamp(value, lo, hi):
+    """Produces a tensor of values from `value` clamped to the range ``[lo, hi]``.
+
+    :rtype: Tensor or Batch of the type that is calculated based on the type promotion rules.
+    """
+    return _arithm_op("clamp", value, lo, hi)