mirror of
https://github.com/saymrwulf/pytorch.git
synced 2026-05-15 21:00:47 +00:00
b9b160d86f
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/16115 Reviewed By: bddppq Differential Revision: D13717049 fbshipit-source-id: fb1d690183a932a1fa1a2d235f3219520f51620a
85 lines
2.2 KiB
C++
85 lines
2.2 KiB
C++
#include <torch/extension.h>
|
|
|
|
#include <ATen/CPUFloatType.h>
|
|
#include <ATen/Type.h>
|
|
#include <ATen/core/VariableHooksInterface.h>
|
|
#include <ATen/detail/ComplexHooksInterface.h>
|
|
|
|
#include "ATen/Allocator.h"
|
|
#include "ATen/CPUGenerator.h"
|
|
#include "ATen/DeviceGuard.h"
|
|
#include "ATen/NativeFunctions.h"
|
|
#include "ATen/Utils.h"
|
|
#include "ATen/WrapDimUtils.h"
|
|
#include "c10/util/Half.h"
|
|
#include <c10/core/TensorImpl.h>
|
|
#include <c10/core/UndefinedTensorImpl.h>
|
|
#include <c10/util/Optional.h>
|
|
|
|
#include <cstddef>
|
|
#include <functional>
|
|
#include <memory>
|
|
#include <utility>
|
|
|
|
#include "ATen/Config.h"
|
|
|
|
namespace at {
|
|
|
|
struct CPUComplexFloatType : public at::CPUTypeDefault {
|
|
CPUComplexFloatType()
|
|
: CPUTypeDefault(
|
|
CPUTensorId(),
|
|
/*is_variable=*/false,
|
|
/*is_undefined=*/false) {}
|
|
|
|
ScalarType scalarType() const override;
|
|
caffe2::TypeMeta typeMeta() const override;
|
|
Backend backend() const override;
|
|
const char* toString() const override;
|
|
size_t elementSizeInBytes() const override;
|
|
TypeID ID() const override;
|
|
|
|
Tensor empty(IntList size, const TensorOptions & options) const override {
|
|
// Delegate to the appropriate cpu tensor factory
|
|
const DeviceGuard device_guard(options.device());
|
|
return at::native::empty_cpu(/* actuals */ size, options);
|
|
}
|
|
};
|
|
|
|
struct ComplexHooks : public at::ComplexHooksInterface {
|
|
ComplexHooks(ComplexHooksArgs) {}
|
|
void registerComplexTypes(Context* context) const override {
|
|
context->registerType(
|
|
Backend::CPU, ScalarType::ComplexFloat, new CPUComplexFloatType());
|
|
}
|
|
};
|
|
|
|
ScalarType CPUComplexFloatType::scalarType() const {
|
|
return ScalarType::ComplexFloat;
|
|
}
|
|
|
|
caffe2::TypeMeta CPUComplexFloatType::typeMeta() const {
|
|
return scalarTypeToTypeMeta(ScalarType::ComplexFloat);
|
|
}
|
|
|
|
Backend CPUComplexFloatType::backend() const {
|
|
return Backend::CPU;
|
|
}
|
|
|
|
const char* CPUComplexFloatType::toString() const {
|
|
return "CPUComplexFloatType";
|
|
}
|
|
|
|
TypeID CPUComplexFloatType::ID() const {
|
|
return TypeID::CPUComplexFloat;
|
|
}
|
|
|
|
size_t CPUComplexFloatType::elementSizeInBytes() const {
|
|
return sizeof(float);
|
|
}
|
|
|
|
REGISTER_COMPLEX_HOOKS(ComplexHooks);
|
|
|
|
} // namespace at
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { }
|