ScalarOps.h 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. #pragma once
  2. #include <c10/core/Scalar.h>
  3. #include <ATen/Tensor.h>
  4. #ifndef AT_PER_OPERATOR_HEADERS
  5. #include <ATen/Functions.h>
  6. #else
  7. #include <ATen/ops/scalar_tensor.h>
  8. #endif
  9. namespace at {
  10. namespace detail {
  11. // When filling a number to 1-element CPU tensor, we want to skip
  12. // everything but manipulate data ptr directly.
  13. // Ideally this fast pass should be implemented in TensorIterator,
  14. // but we also want to skip compute_types which in not avoidable
  15. // in TensorIterator for now.
  16. Tensor& scalar_fill(Tensor& self, const Scalar& value);
  17. TORCH_API Tensor scalar_tensor_static(const Scalar& s, c10::optional<ScalarType> dtype_opt, c10::optional<Device> device_opt);
  18. } // namespace detail
  19. } // namespace at
  20. // This is in the c10 namespace because we use ADL to find the functions in it.
  21. namespace c10 {
  22. // FIXME: this should be (and was) Scalar::toTensor, but there is currently no way
  23. // to implement this without going through Derived Types (which are not part of core).
  24. inline at::Tensor scalar_to_tensor(const Scalar& s, const Device device = at::kCPU) {
  25. // This is the fast track we have for CPU scalar tensors.
  26. if (device == at::kCPU) {
  27. if (s.isFloatingPoint()) {
  28. return at::detail::scalar_tensor_static(s, at::kDouble, at::kCPU);
  29. } else if (s.isComplex()) {
  30. return at::detail::scalar_tensor_static(s, at::kComplexDouble, at::kCPU);
  31. } else if (s.isBoolean()) {
  32. return at::detail::scalar_tensor_static(s, at::kBool, at::kCPU);
  33. } else {
  34. AT_ASSERT(s.isIntegral(false));
  35. return at::detail::scalar_tensor_static(s, at::kLong, at::kCPU);
  36. }
  37. }
  38. if (s.isFloatingPoint()) {
  39. return at::scalar_tensor(s, at::device(device).dtype(at::kDouble));
  40. } else if (s.isBoolean()) {
  41. return at::scalar_tensor(s, at::device(device).dtype(at::kBool));
  42. } else if (s.isComplex()) {
  43. return at::scalar_tensor(s, at::device(device).dtype(at::kComplexDouble));
  44. } else {
  45. AT_ASSERT(s.isIntegral(false));
  46. return at::scalar_tensor(s, at::device(device).dtype(at::kLong));
  47. }
  48. }
  49. } // namespace c10
  50. namespace at {
  51. namespace native {
  52. inline Tensor wrapped_scalar_tensor(const Scalar& scalar, const Device device = at::kCPU) {
  53. auto tensor = scalar_to_tensor(scalar, device);
  54. tensor.unsafeGetTensorImpl()->set_wrapped_number(true);
  55. return tensor;
  56. }
  57. } // namespace native
  58. } // namsepace at