bool BroadcastDimsProductEqual(Value input, Value output,
                               size_t agg_start_idx) {
  ArrayRef<int64_t> input_shape =
      mlir::cast<ShapedType>(input.getType()).getShape();
  ArrayRef<int64_t> output_shape =
      mlir::cast<ShapedType>(output.getType()).getShape();

  int64_t agg_value = 1;
  for (size_t i = agg_start_idx; i < input_shape.size() - 1; ++i) {
    agg_value *= input_shape[i];
  }

// Return true if the product of dimension values of a subsection of the tensor
// is equal to the non-contracting dimension after a reshape
class BroadcastDimsProductEqual<int agg_start_idx> : Constraint<CPred<
  "TFL::BroadcastDimsProductEqual($0, $1, "# agg_start_idx #")">>;

// Returns true if the dimensions of a subsection of two tensors is equal
class AreTensorSubSectionShapesEqual<int skip_first, int skip_last> : Constraint<CPred<