You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
52 lines
1.7 KiB
52 lines
1.7 KiB
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
|
|
#pragma once
|
|
|
|
#include <ATen/native/DispatchStub.h>
|
|
#include <ATen/OpMathType.h>
|
|
#include <ATen/TensorIterator.h>
|
|
#include <c10/core/Scalar.h>
|
|
|
|
namespace at::native {
|
|
|
|
template <typename scalar_t>
|
|
C10_HOST_DEVICE C10_ALWAYS_INLINE bool is_lerp_weight_small(scalar_t weight) {
|
|
return std::abs(weight) < scalar_t(0.5);
|
|
}
|
|
template <typename scalar_t>
|
|
C10_HOST_DEVICE C10_ALWAYS_INLINE bool is_lerp_weight_small(c10::complex<scalar_t> weight) {
|
|
// Avoid the sqrt in abs(weight)
|
|
return (weight.real() * weight.real() + weight.imag() * weight.imag()) < scalar_t(0.25);
|
|
}
|
|
|
|
template <typename scalar_t, typename weight_t>
|
|
C10_HOST_DEVICE C10_ALWAYS_INLINE scalar_t lerp(scalar_t self_, scalar_t end_, weight_t weight_) {
|
|
using opmath_t = at::opmath_type<scalar_t>;
|
|
using opmath_weight_t = at::opmath_type<weight_t>;
|
|
|
|
opmath_t self = self_;
|
|
opmath_t end = end_;
|
|
opmath_weight_t weight = weight_;
|
|
|
|
// Conditional for better numeric. This has been discussed in
|
|
// https://github.com/pytorch/pytorch/pull/18871
|
|
return is_lerp_weight_small(weight)
|
|
? self + weight * (end - self)
|
|
: end - (end - self) * (opmath_t(1) - weight);
|
|
}
|
|
|
|
using lerp_fn_scalar = void (*)(
|
|
at::TensorIteratorBase& iter,
|
|
const Scalar& weight);
|
|
|
|
using lerp_fn_tensor = void (*)(
|
|
at::TensorIteratorBase& iter);
|
|
|
|
DECLARE_DISPATCH(lerp_fn_scalar, lerp_kernel_scalar_weight)
|
|
DECLARE_DISPATCH(lerp_fn_tensor, lerp_kernel_tensor_weight)
|
|
|
|
} // namespace at::native
|
|
|
|
#else
|
|
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
|
|
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
|