#include "caffe2/operators/segment_reduction_op.h"

namespace caffe2 {

OpSchema::Cost CostInferenceForSparseLengths(
    const OperatorDef& def,
    const vector<TensorShape>& inputs,
    bool use_weight) {
  int min_num_of_inputs = 3 + use_weight;
  CAFFE_ENFORCE_GE(
      inputs.size(),
      min_num_of_inputs,
      def.type() + " requires at least " + c10::to_string(min_num_of_inputs));

  const TensorShape data = inputs[0];
  const TensorShape indices = inputs[1 + use_weight];
  const TensorShape lengths = inputs[2 + use_weight];

  OpSchema::Cost c;
  CAFFE_ENFORCE_GT(data.dims_size(), 0, "data requires at least 1 dimension");
  uint64_t N = data.dims(0);
  if (N == 0) {
    return c;
  }
  uint64_t D = nElemFromDim(data, 1);
  CAFFE_ENFORCE_GT(
      lengths.dims_size(), 0, "lengths requires at least 1 dimension");
  uint64_t M = lengths.dims(0);
  uint64_t indices_size = nElemFromDim(indices);

  c.flops = indices_size * D;
  c.bytes_read = indices_size *
          (D * sizeof(data.data_type()) + sizeof(indices.data_type())) +
      M * sizeof(lengths.data_type());
  c.params_bytes = N * D * sizeof(data.data_type());
  if (use_weight) {
    const TensorShape weights = inputs[1];
    c.flops += indices_size * D;
    c.bytes_read += indices_size * sizeof(weights.data_type());
  }

  return c;
}

// registering 5 input gradient with main output
// gradient of SparseLengthsWeightedSum
OPERATOR_SCHEMA(SparseLengthsIndicesInGradientWeightedSumWithMainInputGradient)
    .NumInputs(5)
    .NumOutputs(2);
REGISTER_CPU_OPERATOR(
    SparseLengthsIndicesInGradientWeightedSumWithMainInputGradient,
    AbstractLengthsWithMainInputGradientOp<
        float,
        float,
        int,
        CPUContext,
        WeightedSumReducerDef::template ReducerGradient<float, CPUContext>,
        true /*SparseFused*/,
        true /*GradientNeedIndices*/>);

// registering 4 input version
OPERATOR_SCHEMA(SparseLengthsIndicesInGradientWeightedSumGradient)
    .NumInputs(4)
    .NumOutputs(1);
REGISTER_CPU_OPERATOR(
    SparseLengthsIndicesInGradientWeightedSumGradient,
    AbstractLengthsGradientOp<
        float,
        int,
        CPUContext,
        WeightedSumReducerDef::template ReducerGradient<float, CPUContext>,
        true /*GradientNeedIndices*/>);

// registering 3 input version
// gradient of SparseLengthsSum
OPERATOR_SCHEMA(SparseLengthsIndicesInGradientSumGradient)
    .NumInputs(3)
    .NumOutputs(1);
REGISTER_CPU_OPERATOR(
    SparseLengthsIndicesInGradientSumGradient,
    AbstractLengthsGradientOp<
        float,
        int,
        CPUContext,
        SumReducerDef::template ReducerGradient<float, CPUContext>,
        true /*GradientNeedIndices*/>);
// gradient of LengthsSum
OPERATOR_SCHEMA(LengthsIndicesInGradientSumGradient).NumInputs(3).NumOutputs(1);
REGISTER_CPU_OPERATOR(
    LengthsIndicesInGradientSumGradient,
    AbstractLengthsGradientOp<
        float,
        int,
        CPUContext,
        SumReducerDef::template ReducerGradient<float, CPUContext>,
        true /*GradientNeedIndices*/>);

// registering 3 input version
// gradient of SparseLengthsMean
OPERATOR_SCHEMA(SparseLengthsIndicesInGradientMeanGradient)
    .NumInputs(3)
    .NumOutputs(1);
REGISTER_CPU_OPERATOR(
    SparseLengthsIndicesInGradientMeanGradient,
    AbstractLengthsGradientOp<
        float,
        int,
        CPUContext,
        MeanReducerDef::template ReducerGradient<float, CPUContext>,
        true /*GradientNeedIndices*/>);
// gradient of LengthsMean
OPERATOR_SCHEMA(LengthsIndicesInGradientMeanGradient)
    .NumInputs(3)
    .NumOutputs(1);
REGISTER_CPU_OPERATOR(
    LengthsIndicesInGradientMeanGradient,
    AbstractLengthsGradientOp<
        float,
        int,
        CPUContext,
        MeanReducerDef::template ReducerGradient<float, CPUContext>,
        true /*GradientNeedIndices*/>);

namespace {

static const char* kLengthsMaxExtra = R"DOC(
The *LengthsMax* op takes two inputs *DATA* and *LENGTHS*, and produces a single output *OUTPUT*. The op finds the maximum value in each of the segments of *DATA*, where segments are defined by their lengths.
For example, if $DATA = [2,4,3,1,2,10]$ and $LENGTHS = [2,3,1]$ then $OUTPUT = [max([2,4]), max([3,1,2]), max([10])] = [4,3,10]$.

Github Link:
- https://github.com/caffe2/caffe2/blob/master/caffe2/operators/segment_reduction_op.cc

<details>

<summary> <b>Example</b> </summary>

**Code**

```

workspace.ResetWorkspace()

op = core.CreateOperator(
    "LengthsMax",
    ["DATA", "LENGTHS"],
    ["OUTPUT"],
)

workspace.FeedBlob("DATA", np.array([2,4,3,1,2,10]).astype(np.float32))
print("DATA:\n", workspace.FetchBlob("DATA"))

workspace.FeedBlob("LENGTHS", np.array([2,3,1]).astype(np.int32))
print("LENGTHS:\n", workspace.FetchBlob("LENGTHS"))

workspace.RunOperatorOnce(op)
print("OUTPUT: \n", workspace.FetchBlob("OUTPUT"))

```

**Result**

```

DATA:
 [ 2.  4.  3.  1.  2. 10.]
LENGTHS:
 [2 3 1]
OUTPUT:
 [ 4.  3. 10.]

```

</details>

)DOC";

static const char* kLengthsMeanExtra = R"DOC(
The *LengthsMean* op takes two inputs *DATA* and *LENGTHS*, and produces a single output *OUTPUT*. The op finds the mean value in each of the segments of *DATA*, where segments are defined by their lengths.
For example, if $DATA = [2,4,3,1,2,10]$ and $LENGTHS = [2,3,1]$ then $OUTPUT = [mean([2,4]), mean([3,1,2]), mean([10])] = [3,2,10]$.

Github Link:
- https://github.com/caffe2/caffe2/blob/master/caffe2/operators/segment_reduction_op.cc

<details>

<summary> <b>Example</b> </summary>

**Code**

```

workspace.ResetWorkspace()

op = core.CreateOperator(
    "LengthsMean",
    ["DATA", "LENGTHS"],
    ["OUTPUT"],
)

workspace.FeedBlob("DATA", np.array([2,4,3,1,2,10]).astype(np.float32))
print("DATA:\n", workspace.FetchBlob("DATA"))

workspace.FeedBlob("LENGTHS", np.array([2,3,1]).astype(np.int32))
print("LENGTHS:\n", workspace.FetchBlob("LENGTHS"))

workspace.RunOperatorOnce(op)
print("OUTPUT: \n", workspace.FetchBlob("OUTPUT"))

```

**Result**

```

DATA:
 [ 2.  4.  3.  1.  2. 10.]
LENGTHS:
 [2 3 1]
OUTPUT:
 [ 3.  2. 10.]

```

</details>

)DOC";

static const char* kLengthsSumExtra = R"DOC(
The *LengthsSum* op takes two inputs *DATA* and *LENGTHS*, and produces a single output *OUTPUT*. The op finds the sum in each of the segments of *DATA*, where segments are defined by their lengths.
For example, if $DATA = [2,4,3,1,2,10]$ and $LENGTHS = [2,3,1]$ then $OUTPUT = [sum([2,4]), sum([3,1,2]), sum([10])] = [6,6,10]$.

Github Link:
- https://github.com/caffe2/caffe2/blob/master/caffe2/operators/segment_reduction_op.cc

<details>

<summary> <b>Example</b> </summary>

**Code**

```

workspace.ResetWorkspace()

op = core.CreateOperator(
    "LengthsSum",
    ["DATA", "LENGTHS"],
    ["OUTPUT"],
)

workspace.FeedBlob("DATA", np.array([2,4,3,1,2,10]).astype(np.float32))
print("DATA:\n", workspace.FetchBlob("DATA"))

workspace.FeedBlob("LENGTHS", np.array([2,3,1]).astype(np.int32))
print("LENGTHS:\n", workspace.FetchBlob("LENGTHS"))

workspace.RunOperatorOnce(op)
print("OUTPUT: \n", workspace.FetchBlob("OUTPUT"))

```

**Result**

```

DATA:
 [ 2.  4.  3.  1.  2. 10.]
LENGTHS:
 [2 3 1]
OUTPUT:
 [ 6.  6. 10.]

```

</details>

)DOC";

static const char* kLengthsWeightedSumExtra = R"DOC(
The *LengthsWeightedSum* op takes three inputs *DATA*, *LENGTHS*, and *SCALARS*, and produces a single output *OUTPUT*. The op finds the weighted sum in each of the segments of *DATA*, where segments are defined by their lengths. Before calculating the sums, the input *DATA* is weighted by the contents of *SCALARS*.
For example, if $DATA = [2,4,3,1,2,10]$, $SCALARS = [8, 2, 1, 4, 1, 0.6]$, and $LENGTHS = [2,3,1]$, then $OUTPUT = [sum([8*2,2*4]), sum([1*3,4*1,1*2]), sum([0.6*10])] = [24,9,6]$.

Github Link:
- https://github.com/caffe2/caffe2/blob/master/caffe2/operators/segment_reduction_op.cc

<details>

<summary> <b>Example</b> </summary>

**Code**

```

workspace.ResetWorkspace()

op = core.CreateOperator(
    "LengthsWeightedSum",
    ["DATA", "SCALARS","LENGTHS"],
    ["OUTPUT"],
)

workspace.FeedBlob("DATA", np.array([2,4,3,1,2,10]).astype(np.float32))
print("DATA:\n", workspace.FetchBlob("DATA"))

workspace.FeedBlob("SCALARS", np.array([8, 2, 1, 4, 1, 0.6]).astype(np.float32))
print("SCALARS:\n", workspace.FetchBlob("SCALARS"))

workspace.FeedBlob("LENGTHS", np.array([2,3,1]).astype(np.int32))
print("LENGTHS:\n", workspace.FetchBlob("LENGTHS"))

workspace.RunOperatorOnce(op)
print("OUTPUT: \n", workspace.FetchBlob("OUTPUT"))

```

**Result**

```

DATA:
 [ 2.  4.  3.  1.  2. 10.]
SCALARS:
 [8.  2.  1.  4.  1.  0.6]
LENGTHS:
 [2 3 1]
OUTPUT:
 [24.  9.  6.]

```

</details>

)DOC";

template <typename Def>
string FormatDoc() {
  string doc = Def::doc;
  c10::ReplaceAll(doc, "{op}", Def::OpDef::name);
  c10::ReplaceAll(doc, "{op_doc}", Def::OpDef::doc);
  if (strcmp(Def::OpDef::name, "Max") == 0) {
    c10::ReplaceAll(doc, "{extra}", kLengthsMaxExtra);
  } else if (strcmp(Def::OpDef::name, "Mean") == 0) {
    c10::ReplaceAll(doc, "{extra}", kLengthsMeanExtra);
  } else if (strcmp(Def::OpDef::name, "Sum") == 0) {
    c10::ReplaceAll(doc, "{extra}", kLengthsSumExtra);
  } else if (strcmp(Def::OpDef::name, "WeightedSum") == 0) {
    c10::ReplaceAll(doc, "{extra}", kLengthsWeightedSumExtra);
  } else {
    c10::ReplaceAll(doc, "{extra}", " ");
  }
  return doc;
}

// Helper function to enforce naming conventions at compile time.
constexpr bool equal(
    char const* lhs,
    char const* rhs1,
    char const* rhs2,
    char const* rhs3 = "") {
  return (*lhs == 0 && *rhs1 == 0 && *rhs2 == 0 && *rhs3 == 0) ||
      (*rhs1 != 0 && *lhs == *rhs1 && equal(lhs + 1, rhs1 + 1, rhs2, rhs3)) ||
      (*rhs1 == 0 && *rhs2 != 0 && *lhs == *rhs2 &&
       equal(lhs + 1, rhs1, rhs2 + 1, rhs3)) ||
      (*rhs1 == 0 && *rhs2 == 0 && *rhs3 != 0 && *lhs == *rhs3 &&
       equal(lhs + 1, rhs1, rhs2, rhs3 + 1));
}

// Helper macro when the main op is defined elsewhere, and we only need to
// define the schema, and the gradient op.
// TODO: enable input fillers
#define REGISTER_SEGMENT_DEF_SCHEMA_GRADIENT_ONLY(                            \
    segment_name, gradient_name, ...)                                         \
  static_assert(                                                              \
      equal(#segment_name, __VA_ARGS__::basename, __VA_ARGS__::OpDef::name),  \
      #segment_name);                                                         \
  static_assert(                                                              \
      equal(                                                                  \
          #gradient_name,                                                     \
          __VA_ARGS__::basename,                                              \
          __VA_ARGS__::OpDef::name,                                           \
          "Gradient"),                                                        \
      #gradient_name);                                                        \
  OPERATOR_SCHEMA(segment_name)                                               \
      .NumInputs(__VA_ARGS__::ForwardOp::kNumInputs)                          \
      .NumOutputs(1)                                                          \
      .DisallowInputFillers()                                                 \
      .SetDoc(FormatDoc<__VA_ARGS__>())                                       \
      .Output(0, "OUTPUT", "Aggregated tensor")                               \
      .FillUsing(__VA_ARGS__::PopulateSchema);                                \
  REGISTER_CPU_OPERATOR_STR(string(#gradient_name), __VA_ARGS__::BackwardOp); \
  OPERATOR_SCHEMA(gradient_name)                                              \
      .NumInputs(__VA_ARGS__::BackwardOp::kNumInputs)                         \
      .NumOutputs(1)                                                          \
      .DisallowInputFillers();                                                \
  REGISTER_GRADIENT_STR(string(#segment_name), __VA_ARGS__::GetGradient)

#define REGISTER_SEGMENT_DEF(segment_name, gradient_name, ...)               \
  static_assert(                                                             \
      equal(#segment_name, __VA_ARGS__::basename, __VA_ARGS__::OpDef::name), \
      #segment_name);                                                        \
  REGISTER_CPU_OPERATOR_STR(string(#segment_name), __VA_ARGS__::ForwardOp);  \
  REGISTER_SEGMENT_DEF_SCHEMA_GRADIENT_ONLY(                                 \
      segment_name, gradient_name, __VA_ARGS__)

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentRangeSum,
    SortedSegmentRangeSumGradient,
    AbstractSortedSegmentRangeDef<float, int, CPUContext, SumRangeReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentRangeLogSumExp,
    SortedSegmentRangeLogSumExpGradient,
    AbstractSortedSegmentRangeDef<
        float,
        int,
        CPUContext,
        LogSumExpRangeReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentRangeLogMeanExp,
    SortedSegmentRangeLogMeanExpGradient,
    AbstractSortedSegmentRangeDef<
        float,
        int,
        CPUContext,
        LogMeanExpRangeReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentRangeMean,
    SortedSegmentRangeMeanGradient,
    AbstractSortedSegmentRangeDef<float, int, CPUContext, MeanRangeReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentRangeMax,
    SortedSegmentRangeMaxGradient,
    AbstractSortedSegmentRangeDef<float, int, CPUContext, MaxRangeReducerDef>);

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentSum,
    SortedSegmentSumGradient,
    AbstractSortedSegmentDef<float, int, CPUContext, SumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SparseSortedSegmentSum,
    SparseSortedSegmentSumGradient,
    AbstractSparseSortedSegmentDef<float, int, CPUContext, SumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    UnsortedSegmentSum,
    UnsortedSegmentSumGradient,
    AbstractUnsortedSegmentDef<float, int, CPUContext, SumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SparseUnsortedSegmentSum,
    SparseUnsortedSegmentSumGradient,
    AbstractSparseUnsortedSegmentDef<float, int, CPUContext, SumReducerDef>);

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    LengthsSum,
    LengthsSumGradient,
    AbstractLengthsDef<float, int, CPUContext, SumReducerDef, true>);

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentMean,
    SortedSegmentMeanGradient,
    AbstractSortedSegmentDef<float, int, CPUContext, MeanReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SparseSortedSegmentMean,
    SparseSortedSegmentMeanGradient,
    AbstractSparseSortedSegmentDef<float, int, CPUContext, MeanReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    UnsortedSegmentMean,
    UnsortedSegmentMeanGradient,
    AbstractUnsortedSegmentDef<float, int, CPUContext, MeanReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SparseUnsortedSegmentMean,
    SparseUnsortedSegmentMeanGradient,
    AbstractSparseUnsortedSegmentDef<float, int, CPUContext, MeanReducerDef>);

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    LengthsMean,
    LengthsMeanGradient,
    AbstractLengthsDef<float, int, CPUContext, MeanReducerDef, true>);

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    ReduceFrontWeightedSum,
    ReduceFrontWeightedSumGradient,
    AbstractReduceFrontDef<float, CPUContext, WeightedSumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SortedSegmentWeightedSum,
    SortedSegmentWeightedSumGradient,
    AbstractSortedSegmentDef<float, int, CPUContext, WeightedSumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SparseSortedSegmentWeightedSum,
    SparseSortedSegmentWeightedSumGradient,
    AbstractSparseSortedSegmentDef<
        float,
        int,
        CPUContext,
        WeightedSumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    UnsortedSegmentWeightedSum,
    UnsortedSegmentWeightedSumGradient,
    AbstractUnsortedSegmentDef<float, int, CPUContext, WeightedSumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    SparseUnsortedSegmentWeightedSum,
    SparseUnsortedSegmentWeightedSumGradient,
    AbstractSparseUnsortedSegmentDef<
        float,
        int,
        CPUContext,
        WeightedSumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_SEGMENT_DEF(
    LengthsWeightedSum,
    LengthsWeightedSumGradient,
    AbstractLengthsDef<float, int, CPUContext, WeightedSumReducerDef, false>);

// Auxiliary output gradients are currently implemented only for Lengths version
#define REGISTER_GRADIENT_WITH_MAIN_INPUT(gradient_name, ...)        \
  static_assert(                                                     \
      equal(                                                         \
          #gradient_name,                                            \
          __VA_ARGS__::basename,                                     \
          __VA_ARGS__::OpDef::name,                                  \
          "WithMainInputGradient"),                                  \
      #gradient_name);                                               \
  REGISTER_CPU_OPERATOR_STR(                                         \
      string(#gradient_name), __VA_ARGS__::WithMainInputBackwardOp); \
  OPERATOR_SCHEMA(gradient_name)                                     \
      .NumInputs(__VA_ARGS__::WithMainInputBackwardOp::kNumInputs)   \
      .NumOutputs(1, INT_MAX)

// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_GRADIENT_WITH_MAIN_INPUT(
    LengthsWeightedSumWithMainInputGradient,
    AbstractLengthsDef<float, int, CPUContext, WeightedSumReducerDef>);
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables,clang-diagnostic-unused-function)
REGISTER_GRADIENT_WITH_MAIN_INPUT(
    SparseLengthsWeightedSumWithMainInputGradient,
    AbstractSparseLengthsDef<float, int, CPUContext, WeightedSumReducerDef>);
} // namespace

#define REGISTER_GRADIENT_WITH_MAIN_INPUT_AND_FORWARD_OUTPUT(               \
    gradient_name, ...)                                                     \
  static_assert(                                                            \
      equal(                                                                \
          #gradient_name,                                                   \
          __VA_ARGS__::basename,                                            \
          __VA_ARGS__::OpDef::name,                                         \
          "WithMainInputAndForwardOutputGradient"),                         \
      #gradient_name);                                                      \
  REGISTER_CPU_OPERATOR_STR(                                                \
      string(#gradient_name),                                               \
      __VA_ARGS__::WithMainInputAndForwardOutputBackwardOp);                \
  OPERATOR_SCHEMA(gradient_name)                                            \
      .NumInputs(                                                           \
          __VA_ARGS__::WithMainInputAndForwardOutputBackwardOp::kNumInputs) \
      .NumOutputs(1, INT_MAX)

#define REGISTER_SEGMENT_DEF_MAIN_INPUT_AND_FORWARD_OUTPUT_GRADIENT(         \
    segment_name, gradient_name, ...)                                        \
  static_assert(                                                             \
      equal(#segment_name, __VA_ARGS__::basename, __VA_ARGS__::OpDef::name), \
      #segment_name);                                                        \
  OPERATOR_SCHEMA(segment_name)                                              \
      .NumInputs(__VA_ARGS__::ForwardOp::kNumInputs)                         \
      .NumOutputs(1)                                                         \
      .SetDoc(FormatDoc<__VA_ARGS__>())                                      \
      .Output(0, "OUTPUT", "Aggregated tensor")                              \
      .FillUsing(__VA_ARGS__::PopulateSchema);                               \
  REGISTER_GRADIENT_WITH_MAIN_INPUT_AND_FORWARD_OUTPUT(                      \
      gradient_name, __VA_ARGS__);                                           \
  REGISTER_GRADIENT_STR(string(#segment_name), __VA_ARGS__::GetGradient)

// This implements and registers a length op with a gradient which requires
// the main input as well as the output of the forward output.
#define REGISTER_LENGTHS_OPS_MAIN_INPUT_AND_FORWARD_OUTPUT_GRADIENT(         \
    segment_name, gradient_name, ...)                                        \
  static_assert(                                                             \
      equal(#segment_name, __VA_ARGS__::basename, __VA_ARGS__::OpDef::name), \
      #segment_name);                                                        \
  REGISTER_CPU_OPERATOR_STR(string(#segment_name), __VA_ARGS__::ForwardOp);  \
  REGISTER_SEGMENT_DEF_MAIN_INPUT_AND_FORWARD_OUTPUT_GRADIENT(               \
      segment_name, gradient_name, __VA_ARGS__)

REGISTER_LENGTHS_OPS_MAIN_INPUT_AND_FORWARD_OUTPUT_GRADIENT(
    LengthsMax,
    LengthsMaxWithMainInputAndForwardOutputGradient,
    AbstractLengthsDef<float, int, CPUContext, MaxReducerDef>);
} // namespace caffe2

// Macro doesn't like comma
using LengthsSumCPUOp = caffe2::AbstractLengthsDef<
    float,
    int,
    caffe2::CPUContext,
    caffe2::SumReducerDef,
    true>::ForwardOp;
using LengthsMeanCPUOp = caffe2::AbstractLengthsDef<
    float,
    int,
    caffe2::CPUContext,
    caffe2::MeanReducerDef,
    true>::ForwardOp;
using LengthsMaxCPUOp = caffe2::AbstractLengthsDef<
    float,
    int,
    caffe2::CPUContext,
    caffe2::MaxReducerDef,
    true>::ForwardOp;

C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
    LengthsSum,
    "_caffe2::LengthsSum(Tensor data, Tensor lengths) -> Tensor",
    LengthsSumCPUOp);
C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
    LengthsMean,
    "_caffe2::LengthsMean(Tensor data, Tensor lengths) -> Tensor",
    LengthsMeanCPUOp);
C10_EXPORT_CAFFE2_OP_TO_C10_CPU(
    LengthsMax,
    "_caffe2::LengthsMax(Tensor data, Tensor lengths) -> Tensor",
    LengthsMaxCPUOp);
