diff --git a/caffe2/operators/gather_ranges_to_dense_op.cc b/caffe2/operators/gather_ranges_to_dense_op.cc index 10396aafc97..aa31ef12b36 100644 --- a/caffe2/operators/gather_ranges_to_dense_op.cc +++ b/caffe2/operators/gather_ranges_to_dense_op.cc @@ -104,3 +104,11 @@ NO_GRADIENT(GatherRangesToDense); } // namespace } // namespace caffe2 + +using GatherRangesToDenseCPUOp = + caffe2::GatherRangesToDenseOp; + +C10_EXPORT_CAFFE2_OP_TO_C10_CPU( + GatherRangesToDense, + "_caffe2::GatherRangesToDense(Tensor data, Tensor ranges, Tensor? key, int[] lengths, int min_observation, float max_mismatched_ratio, float max_empty_ratio) -> Tensor[] outputs", + GatherRangesToDenseCPUOp); diff --git a/caffe2/operators/gather_ranges_to_dense_op.h b/caffe2/operators/gather_ranges_to_dense_op.h index c1dd5a52700..217a61b2512 100644 --- a/caffe2/operators/gather_ranges_to_dense_op.h +++ b/caffe2/operators/gather_ranges_to_dense_op.h @@ -5,6 +5,7 @@ #include "caffe2/core/common_omp.h" #include "caffe2/core/context.h" +#include "caffe2/core/export_caffe2_op_to_c10.h" #include "caffe2/core/logging.h" #include "caffe2/core/operator.h" #include "caffe2/core/types.h" @@ -15,6 +16,8 @@ #include #include +C10_DECLARE_EXPORT_CAFFE2_OP_TO_C10(GatherRangesToDense); + namespace caffe2 { template class GatherRangesToDenseOp final : public Operator { diff --git a/caffe2/python/operator_test/torch_integration_test.py b/caffe2/python/operator_test/torch_integration_test.py index 55f26a89987..9bec6476424 100644 --- a/caffe2/python/operator_test/torch_integration_test.py +++ b/caffe2/python/operator_test/torch_integration_test.py @@ -875,6 +875,47 @@ class TorchIntegration(hu.HypothesisTestCase): ) torch.testing.assert_allclose(expected_output, actual_output.cpu()) + def test_gather_ranges_to_dense_op(self): + data = np.array([1, 2, 3, 4, 5, 6, 7, 8]) + ranges = np.array([[[2, 4]], [[0, 0]]]) + key = np.array([0, 1, 3, 2, 1, 0, 1, 0]) + lengths = np.array([4]) + min_observation = 2 + max_mismatched_ratio = 0.5 + max_empty_ratio = 1.0 + + outputs_name = ["X_{}".format(i) for i in range(len(lengths))] + ref_op = core.CreateOperator( + "GatherRangesToDense", + ["data", "ranges", "key"], + outputs_name, + lengths=lengths, + min_observation=min_observation, + max_mismatched_ratio=max_mismatched_ratio, + max_empty_ratio=max_empty_ratio, + ) + workspace.FeedBlob("data", data) + workspace.FeedBlob("ranges", ranges) + workspace.FeedBlob("key", key) + workspace.RunOperatorOnce(ref_op) + ref_outputs = [] + for output_name in outputs_name: + ref_outputs.append(workspace.FetchBlob(output_name)) + + outputs = torch.ops._caffe2.GatherRangesToDense( + torch.from_numpy(data), + torch.from_numpy(ranges), + torch.from_numpy(key), + lengths=lengths, + min_observation=min_observation, + max_mismatched_ratio=max_mismatched_ratio, + max_empty_ratio=max_empty_ratio, + ) + + self.assertEqual(len(ref_outputs), len(outputs)) + for i in range(0, len(ref_outputs)): + np.testing.assert_array_almost_equal(ref_outputs[i], outputs[i].numpy()) + @given(lengths_0=st.integers(1, 10), lengths_1=st.integers(1, 10)) @settings(deadline=1000) def test_merge_id_lists(self, lengths_0, lengths_1):