mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
508f676c50
Summary: Codemod generated with clangr shard mode, 50 files per diff, clangr code(ndim()->dim()): diffusion/FBS/browse/master/fbcode/caffe2/caffe2/fb/codemods/TensorMethodRename.cpp Reviewed By: salexspb Differential Revision: D12935787 fbshipit-source-id: 303d71d3eb050789af2ab9575e5dcc48f6037086
46 lines
1.2 KiB
C++
46 lines
1.2 KiB
C++
#include <iostream>
|
|
|
|
#include "caffe2/core/flags.h"
|
|
#include "caffe2/operators/utility_ops.h"
|
|
#include <gtest/gtest.h>
|
|
|
|
C10_DECLARE_string(caffe_test_root);
|
|
|
|
namespace caffe2 {
|
|
|
|
static void AddConstInput(
|
|
const vector<int64_t>& shape,
|
|
const float value,
|
|
const string& name,
|
|
Workspace* ws) {
|
|
DeviceOption option;
|
|
CPUContext context(option);
|
|
Blob* blob = ws->CreateBlob(name);
|
|
auto* tensor = BlobGetMutableTensor(blob, CPU);
|
|
tensor->Resize(shape);
|
|
math::Set<float, CPUContext>(
|
|
tensor->numel(), value, tensor->template mutable_data<float>(), &context);
|
|
return;
|
|
}
|
|
|
|
TEST(UtilityOpTest, testReshapeWithScalar) {
|
|
Workspace ws;
|
|
OperatorDef def;
|
|
def.set_name("test_reshape");
|
|
def.set_type("Reshape");
|
|
def.add_input("X");
|
|
def.add_output("XNew");
|
|
def.add_output("OldShape");
|
|
def.add_arg()->CopyFrom(MakeArgument("shape", vector<int64_t>{1}));
|
|
AddConstInput(vector<int64_t>(), 3.14, "X", &ws);
|
|
// execute the op
|
|
unique_ptr<OperatorBase> op(CreateOperator(def, &ws));
|
|
EXPECT_TRUE(op->Run());
|
|
Blob* XNew = ws.GetBlob("XNew");
|
|
const TensorCPU& XNewTensor = XNew->Get<Tensor>();
|
|
EXPECT_EQ(1, XNewTensor.dim());
|
|
EXPECT_EQ(1, XNewTensor.numel());
|
|
}
|
|
|
|
} // namespace caffe2
|