Status ConcatShapeHelper(InferenceContext* c, int start_value_index,
                         int end_value_index, int dim_index) {
  ShapeHandle unused;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(dim_index), 0, &unused));
  const Tensor* concat_dim_t = c->input_tensor(dim_index);
  if (concat_dim_t == nullptr) {
    // Return an unknown shape with same rank as inputs, or an unknown rank
    // if no input's rank is known.

    // Find rank.
    int32_t rank = InferenceContext::kUnknownRank;
    for (int i = start_value_index; i < end_value_index; ++i) {
      if (rank == InferenceContext::kUnknownRank) rank = c->Rank(c->input(i));
      if (rank != InferenceContext::kUnknownRank) {
        break;
      }
    }
    if (rank == InferenceContext::kUnknownRank) {
      c->set_output(0, c->UnknownShape());
      return Status::OK();
    } else if (rank == 0) {
      return errors::InvalidArgument(
          "Can't concatenate scalars (use tf.stack instead)");
    } else {
      for (int i = start_value_index; i < end_value_index; ++i) {
        // Check that all the inputs are of the correct rank.
        TF_RETURN_IF_ERROR(c->WithRank(c->input(i), rank, &unused));
      }
    }
    // Build result of <rank> different unknown dims.
    std::vector<DimensionHandle> dims;
    dims.reserve(rank);
    for (int i = 0; i < rank; ++i) dims.push_back(c->UnknownDim());
    c->set_output(0, c->MakeShape(dims));
    return Status::OK();
  }

  // Merge all the non-concat dims, and sum the concat dim to make an output
  // shape.
  int64_t concat_dim;
  if (concat_dim_t->dtype() == DT_INT32) {
    concat_dim = static_cast<int64_t>(concat_dim_t->flat<int32>()(0));
  } else {
    concat_dim = concat_dim_t->flat<int64_t>()(0);
  }

  // Minimum required number of dimensions.
  const int64 min_rank = concat_dim < 0 ? -concat_dim : concat_dim + 1;

  ShapeHandle output_before;
  ShapeHandle output_after;

  ShapeHandle input = c->input(end_value_index - 1);
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
  TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &output_before));
  DimensionHandle output_middle = c->Dim(input, concat_dim);
  if (concat_dim == -1) {
    output_after = c->Scalar();  // no dimensions.
  } else {
    TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &output_after));
  }

  for (int i = end_value_index - 2; i >= start_value_index; --i) {
    ShapeHandle before;
    ShapeHandle after;
    input = c->input(i);
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(input, min_rank, &input));
    TF_RETURN_IF_ERROR(c->Subshape(input, 0, concat_dim, &before));
    DimensionHandle middle = c->Dim(input, concat_dim);
    if (concat_dim == -1) {
      after = c->Scalar();
    } else {
      TF_RETURN_IF_ERROR(c->Subshape(input, concat_dim + 1, &after));
    }

    TF_RETURN_IF_ERROR(c->Merge(before, output_before, &output_before));
    TF_RETURN_IF_ERROR(c->Add(output_middle, middle, &output_middle));
    TF_RETURN_IF_ERROR(c->Merge(after, output_after, &output_after));
  }

  ShapeHandle s;
  TF_RETURN_IF_ERROR(
      c->Concatenate(output_before, c->Vector(output_middle), &s));
  TF_RETURN_IF_ERROR(c->Concatenate(s, output_after, &s));
  c->set_output(0, s);
  return Status::OK();
}


Fix Segfault in Concat V2 shape function.

PiperOrigin-RevId: 412120654
Change-Id: I3ff915faea694f9ad8b00024e9af2de9909011be

Status Conv2DBackpropInputShape(shape_inference::InferenceContext* c) {
  string data_format_str;
  if (!c->GetAttr("data_format", &data_format_str).ok()) {
    data_format_str = "NHWC";
  }
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }

  // For the rest of this function, output_grad_* describes out_backprop and
  // input_grad_* describes in_backprop.
  ShapeHandle output_grad_shape = c->input(2);
  TF_RETURN_IF_ERROR(c->WithRank(output_grad_shape, 4, &output_grad_shape));
  ShapeHandle filter_shape = c->input(1);
  TF_RETURN_IF_ERROR(c->WithRank(filter_shape, 4, &filter_shape));

  DimensionHandle batch_size_dim;
  DimensionHandle output_grad_depth_dim;
  gtl::InlinedVector<DimensionHandle, 2> output_grad_spatial_dims(2);
  TF_RETURN_IF_ERROR(DimensionsFromShape(
      output_grad_shape, data_format, &batch_size_dim,
      absl::MakeSpan(output_grad_spatial_dims), &output_grad_depth_dim, c));
  DimensionHandle unused;
  TF_RETURN_IF_ERROR(
      c->Merge(output_grad_depth_dim, c->Dim(filter_shape, 3), &unused));

  ShapeHandle specified_input_grad_shape;
  TF_RETURN_IF_ERROR(
      c->MakeShapeFromShapeTensor(0, &specified_input_grad_shape));
  if (c->Rank(specified_input_grad_shape) == InferenceContext::kUnknownRank) {
    TF_RETURN_IF_ERROR(c->WithRank(specified_input_grad_shape, 4,
                                   &specified_input_grad_shape));
  }

  // input_grad_depth_dim doesn't equal c->Dim(filter_shape,2) when the number
  // of groups is larger than 1. If input_sizes is a 4D shape, we collect
  // input_grad_depth_dim from input_sizes; otherwise we compute it as
  // c->Dim(filter_shape,2).
  DimensionHandle input_grad_depth_dim;
  gtl::InlinedVector<DimensionHandle, 2> specified_input_grad_spatial_dims(2);
  int specified_input_grad_rank = c->Rank(specified_input_grad_shape);
  if (specified_input_grad_rank == 4) {
    DimensionHandle specified_batch_size_dim;
    TF_RETURN_IF_ERROR(DimensionsFromShape(
        specified_input_grad_shape, data_format, &specified_batch_size_dim,
        absl::MakeSpan(specified_input_grad_spatial_dims),
        &input_grad_depth_dim, c));
    TF_RETURN_IF_ERROR(
        c->Merge(specified_batch_size_dim, batch_size_dim, &unused));
  } else if (specified_input_grad_rank == 2) {
    specified_input_grad_spatial_dims[0] =
        c->Dim(specified_input_grad_shape, 0);
    specified_input_grad_spatial_dims[1] =
        c->Dim(specified_input_grad_shape, 1);
    input_grad_depth_dim = c->Dim(filter_shape, 2);
  } else {
    return errors::InvalidArgument(
        "Conv2DBackpropInput requires input_sizes to contain 4 values or 2 "
        "values, but got: ",
        specified_input_grad_rank);
  }

  ShapeHandle input_grad_shape;
  TF_RETURN_IF_ERROR(ShapeFromDimensions(
      batch_size_dim, specified_input_grad_spatial_dims, input_grad_depth_dim,
      data_format, /*vect_size=*/absl::nullopt, c, &input_grad_shape));
  c->set_output(0, input_grad_shape);
  return Status::OK();
}

Tensorflow is an Open Source Machine Learning Framework. The implementation of shape inference for `ConcatV2` can be used to trigger a denial of service attack via a segfault caused by a type confusion. The `axis` argument is translated into `concat_dim` in the `ConcatShapeHelper` helper function. Then, a value for `min_rank` is computed based on `concat_dim`. This is then used to validate that the `values` tensor has at least the required rank. However, `WithRankAtLeast` receives the lower bound as a 64-bits value and then compares it against the maximum 32-bits integer value that could be represented. Due to the fact that `min_rank` is a 32-bits value and the value of `axis`, the `rank` argument is a negative value, so the error check is bypassed. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.


Status QuantizeV2Shape(InferenceContext* c) {
  int axis = -1;
  Status s = c->GetAttr("axis", &axis);
  if (!s.ok() && s.code() != error::NOT_FOUND) {
    return s;
  }
  if (axis < -1) {
    return errors::InvalidArgument("axis should be at least -1, got ", axis);
  }
  const int minmax_rank = (axis == -1) ? 0 : 1;
  TF_RETURN_IF_ERROR(shape_inference::UnchangedShape(c));
  ShapeHandle minmax;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), minmax_rank, &minmax));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), minmax_rank, &minmax));
  if (axis != -1) {
    ShapeHandle input;
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), axis + 1, &input));
    DimensionHandle depth;
    TF_RETURN_IF_ERROR(
        c->Merge(c->Dim(minmax, 0), c->Dim(input, axis), &depth));
  }
  c->set_output(1, minmax);
  c->set_output(2, minmax);
  return Status::OK();
}


Status DepthwiseConv2DNativeShapeImpl(shape_inference::InferenceContext* c,
                                      bool supports_explicit_padding) {
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &input_shape));
  ShapeHandle filter_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 4, &filter_shape));

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));

  if (strides.size() != 4) {
    return errors::InvalidArgument(
        "DepthwiseConv2D requires the stride attribute to contain 4 values, "
        "but got: ",
        strides.size());
  }

  std::vector<int32> dilations;
  if (!c->GetAttr("dilations", &dilations).ok()) {
    dilations.resize(4, 1);
  }

  if (dilations.size() != 4) {
    return errors::InvalidArgument(
        "DepthwiseConv2D requires the dilations attribute to contain 4 values, "
        "but got: ",
        dilations.size());
  }

  string data_format_str;
  Status s = c->GetAttr("data_format", &data_format_str);
  TensorFormat data_format;
  if (!s.ok() || !FormatFromString(data_format_str, &data_format)) {
    data_format = FORMAT_NHWC;
  }
  int32_t stride_rows;
  int32_t stride_cols;
  int32_t dilation_rows;
  int32_t dilation_cols;
  if (data_format == FORMAT_NCHW) {
    // Canonicalize input shape to NHWC so the shape inference code below can
    // process it.
    input_shape =
        c->MakeShape({{c->Dim(input_shape, 0), c->Dim(input_shape, 2),
                       c->Dim(input_shape, 3), c->Dim(input_shape, 1)}});
    stride_rows = strides[2];
    stride_cols = strides[3];
    dilation_rows = dilations[2];
    dilation_cols = dilations[3];
  } else {
    stride_rows = strides[1];
    stride_cols = strides[2];
    dilation_rows = dilations[1];
    dilation_cols = dilations[2];
  }

  DimensionHandle batch_size_dim = c->Dim(input_shape, 0);
  DimensionHandle in_rows_dim = c->Dim(input_shape, 1);
  DimensionHandle in_cols_dim = c->Dim(input_shape, 2);

  DimensionHandle filter_rows_dim = c->Dim(filter_shape, 0);
  DimensionHandle filter_cols_dim = c->Dim(filter_shape, 1);
  DimensionHandle input_depth = c->Dim(filter_shape, 2);
  DimensionHandle depth_multiplier = c->Dim(filter_shape, 3);

  // Check that the input depths are compatible.
  TF_RETURN_IF_ERROR(
      c->Merge(c->Dim(input_shape, 3), input_depth, &input_depth));

  DimensionHandle output_depth;
  TF_RETURN_IF_ERROR(c->Multiply(input_depth, depth_multiplier, &output_depth));

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  std::vector<int64_t> explicit_paddings;
  if (supports_explicit_padding) {
    Status status = c->GetAttr("explicit_paddings", &explicit_paddings);
    // Use the default value, which is an empty list, if the attribute is not
    // found. Otherwise return the error to the caller.
    if (!status.ok() && !errors::IsNotFound(status)) {
      return status;
    }
    TF_RETURN_IF_ERROR(CheckValidPadding(padding, explicit_paddings,
                                         /*num_dims=*/4, data_format));
  } else {
    DCHECK(padding != Padding::EXPLICIT);
  }

  // TODO(mrry,shlens): Raise an error if the stride would cause
  // information in the input to be ignored. This will require a change
  // in the kernel implementation.
  DimensionHandle output_rows, output_cols;
  int64_t pad_rows_before = -1, pad_rows_after = -1;
  int64_t pad_cols_before = -1, pad_cols_after = -1;
  if (padding == Padding::EXPLICIT) {
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'H',
                             &pad_rows_before, &pad_rows_after);
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'W',
                             &pad_cols_before, &pad_cols_after);
  }
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_rows_dim, filter_rows_dim, dilation_rows, stride_rows, padding,
      pad_rows_before, pad_rows_after, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, in_cols_dim, filter_cols_dim, dilation_cols, stride_cols, padding,
      pad_cols_before, pad_cols_after, &output_cols));

  ShapeHandle output_shape;
  if (data_format == FORMAT_NCHW) {
    output_shape =
        c->MakeShape({batch_size_dim, output_depth, output_rows, output_cols});
  } else {
    output_shape =
        c->MakeShape({batch_size_dim, output_rows, output_cols, output_depth});
  }
  c->set_output(0, output_shape);
  return Status::OK();
}


Status FusedBatchNormGradShape(shape_inference::InferenceContext* c) {
  string data_format_str;
  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format_str));
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }
  const int rank =
      (data_format_str == "NDHWC" || data_format_str == "NCDHW") ? 5 : 4;
  ShapeHandle y_backprop;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &y_backprop));
  ShapeHandle x;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), rank, &x));

  bool is_training;
  TF_RETURN_IF_ERROR(c->GetAttr("is_training", &is_training));

  int channel_dim_index = GetTensorFeatureDimIndex(rank, data_format);
  DimensionHandle channel_dim = c->Dim(y_backprop, channel_dim_index);
  TF_RETURN_IF_ERROR(
      c->Merge(channel_dim, c->Dim(x, channel_dim_index), &channel_dim));

  // covers scale, mean (reserve_space_1), variance (reserve_space_2)
  for (int i = 2; i < 5; ++i) {
    ShapeHandle vec;
    TF_RETURN_IF_ERROR(c->WithRank(c->input(i), 1, &vec));
    TF_RETURN_IF_ERROR(c->Merge(channel_dim, c->Dim(vec, 0), &channel_dim));
  }

  ShapeHandle x_backprop;
  TF_RETURN_IF_ERROR(
      c->ReplaceDim(y_backprop, channel_dim_index, channel_dim, &x_backprop));
  c->set_output(0, x_backprop);
  c->set_output(1, c->Vector(channel_dim));
  c->set_output(2, c->Vector(channel_dim));
  c->set_output(3, c->Vector(0));
  c->set_output(4, c->Vector(0));
  return Status::OK();
}


Status ReduceScatterShape(shape_inference::InferenceContext* c) {
  shape_inference::ShapeHandle in = c->input(0);
  if (!c->RankKnown(in)) {
    // Input shape unknown, so set unknown output shape.
    c->set_output(0, in);
    return Status::OK();
  }

  shape_inference::ShapeHandle group_assignment_shape = c->input(1);
  if (c->Rank(group_assignment_shape) != 2)
    return errors::InvalidArgument(
        "ReduceScatter group_assignment should be rank 2");

  const Tensor* scatter_dimension = c->input_tensor(2);
  if (!scatter_dimension) {
    c->set_output(0, c->UnknownShape());
    return Status::OK();
  }
  int64_t scatter_dim;
  TF_RETURN_IF_ERROR(c->GetScalarFromTensor(scatter_dimension, &scatter_dim));

  std::vector<shape_inference::DimensionHandle> out_dims;
  out_dims.reserve(c->Rank(in));
  for (int i = 0; i < c->Rank(in); ++i) {
    // If the dimension is the scatter_dimension, then divide the dimension
    // by the partition size in the group_assignment.
    if (i == scatter_dim) {
      shape_inference::DimensionHandle dim = c->Dim(in, i);
      shape_inference::DimensionHandle out_dim;
      TF_RETURN_IF_ERROR(c->Divide(dim, c->Dim(group_assignment_shape, 1),
                                   /*evenly_divisible=*/true, &out_dim));
      out_dims.push_back(out_dim);
    } else {
      out_dims.emplace_back(c->Dim(in, i));
    }
  }
  c->set_output(0, c->MakeShape(out_dims));
  return Status::OK();
}


Status SliceShape(InferenceContext* c) {
  ShapeHandle input = c->input(0);
  ShapeHandle begin_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 1, &begin_shape));
  ShapeHandle sizes_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 1, &sizes_shape));

  // Merge to check compatibility of begin and sizes tensors.
  TF_RETURN_IF_ERROR(c->Merge(begin_shape, sizes_shape, &begin_shape));

  DimensionHandle ndims = c->Dim(begin_shape, 0);
  if (c->ValueKnown(ndims)) {
    TF_RETURN_IF_ERROR(c->WithRank(input, c->Value(ndims), &input));
  }

  // NOTE(mrry): Use MakeShapeFromShapeTensor to handle partially-known
  // values, even though the `begin` value does not represent a shape.
  ShapeHandle begin_value;
  TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(1, &begin_value));

  // We check the tensor value here and will only use
  // `MakeShapeFromShapeTensor` when `sizes_value` is null.
  // The reason is that `sizes` might contain -1, which can't
  // be represented (-1 in the ShapeHandle would mean "unknown").
  const Tensor* sizes_value = c->input_tensor(2);

  if (sizes_value != nullptr) {
    TF_RETURN_IF_ERROR(
        c->WithRank(begin_value, sizes_value->NumElements(), &begin_value));
    std::vector<DimensionHandle> dims;
    // If the begin and sizes tensors are available, then
    // we can be precise about the shape of the output.
    if (sizes_value->dtype() == DT_INT64) {
      TF_RETURN_IF_ERROR(
          SliceHelper<int64_t>(c, begin_value, sizes_value, &dims));
    } else {
      TF_RETURN_IF_ERROR(
          SliceHelper<int32>(c, begin_value, sizes_value, &dims));
    }
    c->set_output(0, c->MakeShape(dims));
    return Status::OK();
  } else {
    // In case `sizes` is not available (`sizes_value` is null),
    // we could try to use `MakeShapeFromShapeTensor` here.
    // If sizes contain -1, we will simply consider it as `Unknown`.
    // This is less than ideal but still an improvement of shape inference.
    // The following is an example that returns [None, 1, None] with this
    // code path:
    //   z = tf.zeros((1, 2, 3))
    //   m = tf.slice(z, [0, 0, 0], [tf.constant(1) + 0, 1, -1])
    //   m.get_shape().as_list()
    ShapeHandle sizes_value;
    TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(2, &sizes_value));
    if (c->RankKnown(sizes_value)) {
      TF_RETURN_IF_ERROR(
          c->WithRank(begin_value, c->Rank(sizes_value), &begin_value));
      std::vector<DimensionHandle> dims;
      dims.reserve(c->Rank(sizes_value));
      for (int i = 0; i < c->Rank(sizes_value); ++i) {
        dims.emplace_back(c->Dim(sizes_value, i));
      }
      c->set_output(0, c->MakeShape(dims));
      return Status::OK();
    }
    // We might know the rank of the input.
    if (c->RankKnown(input)) {
      c->set_output(0, c->UnknownShapeOfRank(c->Rank(input)));
      return Status::OK();
    } else {
      return shape_inference::UnknownShape(c);
    }
  }

  return Status::OK();
}


Status ConcatShape(InferenceContext* c, int num_inputs_to_concat) {
  return ConcatShapeHelper(c, 1 /* start_value_index */,
                           1 + num_inputs_to_concat /* end_value_index */,
                           0 /* dim_index */);
}


Status EinsumShape(shape_inference::InferenceContext* c) {
  // We assume that the equation has a valid format. Either (x),(y)->(z)
  // or (x)->(z), where each of (x), (y) and (z) are concatenation of zero or
  // more latin alphabets and contains at most one ellipsis ('...').
  string equation;
  TF_RETURN_IF_ERROR(c->GetAttr("equation", &equation));
  gtl::InlinedVector<string, 2> input_labels;
  string output_labels;
  TF_RETURN_IF_ERROR(
      ParseEinsumEquation(equation, &input_labels, &output_labels));

  if (c->num_inputs() == 0 || c->num_inputs() > 2) {
    return errors::InvalidArgument("Expected either 1 or 2 inputs but got: ",
                                   c->num_inputs());
  }
  const int input_labels_size = input_labels.size();
  if (c->num_inputs() != input_labels_size) {
    return errors::InvalidArgument("Expected ", input_labels.size(),
                                   " inputs for equation ", equation,
                                   " but got: ", c->num_inputs());
  }

  // Validate input subscripts, build the label to dimension mapping and obtain
  // the broadcast shapes that map to ellipsis.
  absl::flat_hash_map<char, DimensionHandle> label_to_dimension;
  gtl::InlinedVector<ShapeHandle, 2> input_bcast_shapes(c->num_inputs());
  for (int i = 0, end = c->num_inputs(); i < end; ++i) {
    bool has_ellipsis = false;
    TF_RETURN_IF_ERROR(ValidateEinsumEllipsis(input_labels[i], &has_ellipsis));
    ShapeHandle input_shape = c->input(i);
    // Validate that the input rank is sufficient for the given number of named
    // labels.
    if (c->RankKnown(input_shape)) {
      if (has_ellipsis) {
        const int num_named_labels =
            static_cast<int>(input_labels[i].size()) - 3;
        TF_RETURN_WITH_CONTEXT_IF_ERROR(
            c->WithRankAtLeast(input_shape, num_named_labels, &input_shape),
            " for ", i, "th input and equation: ", equation);
      } else {
        const int num_named_labels = static_cast<int>(input_labels[i].size());
        TF_RETURN_WITH_CONTEXT_IF_ERROR(
            c->WithRank(input_shape, num_named_labels, &input_shape), " for ",
            i, "th input and equation: ", equation);
      }
    }

    bool seen_ellipsis = false;
    input_bcast_shapes[i] = c->Scalar();
    // Run through the input labels; populate label_to_dimension mapping and
    // compute the broadcast shapes corresponding to the ellipsis (if present).
    for (int label_idx = 0, end = input_labels[i].size(); label_idx < end;
         ++label_idx) {
      const char label = input_labels[i][label_idx];
      // Calculate the input axis that the current label is referring to. After
      // the ellipsis, the axis may be found by using negative indices; i.e the
      // (rank - k)th dimension corresponds to the (num_labels - k)th label.
      const int64_t axis_before_ellipsis = label_idx;
      const int64_t axis_after_ellipsis =
          c->RankKnown(input_shape)
              ? label_idx + c->Rank(input_shape) - input_labels[i].size()
              : -1;

      // Populate the input broadcast shape when we encounter an ellipsis (...).
      if (label == '.') {
        if (!c->RankKnown(input_shape)) {
          input_bcast_shapes[i] = c->UnknownShape();
        } else {
          // The broadcast shape runs till the named label right after the
          // ellipsis, the label with index (label_idx + 3).
          TF_RETURN_IF_ERROR(c->Subshape(input_shape, axis_before_ellipsis,
                                         axis_after_ellipsis + 3,
                                         &input_bcast_shapes[i]));
        }
        label_idx += 2;  // Skip the rest of the ellipsis.
        seen_ellipsis = true;
        continue;
      }
      // Obtain the dimension that the current label corresponds to.
      int64_t axis = seen_ellipsis ? axis_after_ellipsis : axis_before_ellipsis;
      DimensionHandle new_dim = c->RankKnown(input_shape)
                                    ? c->Dim(input_shape, axis)
                                    : c->UnknownDim();
      // If we've seen this label before, make sure previous and current
      // dimensions are compatible.
      if (label_to_dimension.contains(label)) {
        DimensionHandle merged;
        TF_RETURN_IF_ERROR(
            c->Merge(label_to_dimension[label], new_dim, &merged));
        label_to_dimension[label] = merged;
      } else {
        label_to_dimension[label] = new_dim;
      }
    }
  }

  // For two inputs, broadcast the two input broadcast shapes to create the
  // output broadcast shape. For one input, just copy the single broadcast
  // shape.
  ShapeHandle output_bcast_shape;
  if (input_bcast_shapes.size() == 1) {
    output_bcast_shape = input_bcast_shapes[0];
  } else if (input_bcast_shapes.size() == 2) {
    TF_RETURN_IF_ERROR(BroadcastBinaryOpOutputShapeFnHelper(
        c, input_bcast_shapes[0], input_bcast_shapes[1], true,
        &output_bcast_shape));
  }

  bool output_has_ellipsis = false;
  TF_RETURN_IF_ERROR(
      ValidateEinsumEllipsis(output_labels, &output_has_ellipsis));
  if (output_has_ellipsis) {
    // If the output subscript has ellipsis and the output broadcast rank is
    // unknown, then the output shape should have unknown rank.
    if (!c->RankKnown(output_bcast_shape)) {
      c->set_output(0, c->UnknownShape());
      return Status::OK();
    }
  } else {
    // If the output subscripts don't have ellipsis then make sure the output
    // broadcasting shape is empty.
    TF_RETURN_WITH_CONTEXT_IF_ERROR(
        c->WithRankAtMost(output_bcast_shape, 0, &output_bcast_shape),
        " for einsum equation '", equation,
        "' without ellipsis (...) in the output subscripts where input(s) have "
        "non-empty broadcasting shape");
    output_bcast_shape = c->Scalar();
  }

  // Create the output shape from output labels and label_to_dimension mapping.
  std::vector<DimensionHandle> output_dims;
  for (int label_idx = 0, end = output_labels.size(); label_idx < end;
       ++label_idx) {
    const char label = output_labels[label_idx];
    // Append the output_bcast_shape when the ellipsis is encountered.
    if (label == '.') {
      for (int k = 0; k < c->Rank(output_bcast_shape); ++k) {
        output_dims.push_back(c->Dim(output_bcast_shape, k));
      }
      label_idx += 2;  // Skip the rest of the ellipsis.
      continue;
    }
    auto dimension_it = label_to_dimension.find(label);
    if (dimension_it == label_to_dimension.end()) {
      return errors::InvalidArgument(
          "Einsum output subscripts for equation '", equation, "' has label '",
          label, "' which is not present in the input subscripts");
    }
    output_dims.push_back(dimension_it->second);
  }
  c->set_output(0, c->MakeShape(output_dims));
  return Status::OK();
}


Status MatrixDiagPartV2Shape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape, diag_index_shape, unused_shape;
  TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape));
  TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &diag_index_shape));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused_shape));

  const Tensor* diag_index_tensor = c->input_tensor(1);
  if (!c->RankKnown(input_shape) || !c->FullyDefined(diag_index_shape) ||
      diag_index_tensor == nullptr) {
    c->set_output(0, c->UnknownShape());
    return Status::OK();
  }
  int32_t lower_diag_index = 0;
  int32_t upper_diag_index = 0;
  TF_RETURN_IF_ERROR(ReadDiagIndex(c, diag_index_tensor, &lower_diag_index,
                                   &upper_diag_index));
  if (lower_diag_index > upper_diag_index) {
    return errors::InvalidArgument(
        "lower_diag_index is greater than upper_diag_index");
  }

  // Validates lower_diag_index and upper_diag_index.
  const int32_t input_rank = c->Rank(input_shape);
  const int32_t num_rows = c->Value(c->Dim(input_shape, input_rank - 2));
  const int32_t num_cols = c->Value(c->Dim(input_shape, input_rank - 1));
  int32_t max_diag_len = InferenceContext::kUnknownDim;
  if (num_rows != InferenceContext::kUnknownDim &&
      num_cols != InferenceContext::kUnknownDim) {
    if (lower_diag_index != 0 &&  // For when num_rows or num_cols == 0.
        (-num_rows >= lower_diag_index || lower_diag_index >= num_cols)) {
      return errors::InvalidArgument("lower_diag_index is out of bound.");
    }
    if (upper_diag_index != 0 &&  // For when num_rows or num_cols == 0.
        (-num_rows >= upper_diag_index || upper_diag_index >= num_cols)) {
      return errors::InvalidArgument("upper_diag_index is out of bound.");
    }
    max_diag_len = std::min(num_rows + std::min(upper_diag_index, 0),
                            num_cols - std::max(lower_diag_index, 0));
  }

  std::vector<DimensionHandle> dims;
  dims.reserve(input_rank - 2);
  for (int i = 0; i < input_rank - 2; ++i) {
    dims.push_back(c->Dim(input_shape, i));
  }
  if (lower_diag_index < upper_diag_index) {
    dims.push_back(c->MakeDim(upper_diag_index - lower_diag_index + 1));
  }
  dims.push_back(c->MakeDim(max_diag_len));
  c->set_output(0, c->MakeShape(dims));
  return Status::OK();
}


Status ValidateEinsumEllipsis(absl::string_view subscript,
                              bool* found_ellipsis) {
  const int num_periods = absl::c_count(subscript, '.');
  if (num_periods != 0 && num_periods != 3) {
    return errors::InvalidArgument(
        "Expected at most one ellipsis (...), but found ", num_periods,
        " periods (.) in the input subscript: ", subscript);
  }
  if (num_periods == 3 && !absl::StrContains(subscript, "...")) {
    return errors::InvalidArgument(
        "Periods found outside of ellipsis in subscript: ", subscript);
  }
  *found_ellipsis = num_periods > 0;
  return Status::OK();
}


Status ConcatV2Shape(InferenceContext* c) {
  return ConcatShapeHelper(c, 0 /* start_value_index */,
                           c->num_inputs() - 1 /* end_value_index */,
                           c->num_inputs() - 1 /* dim_index */);
}


Status ValidateSparseTensor(InferenceContext* c, ShapeHandle indices_shape,
                            ShapeHandle values_shape, ShapeHandle shape_shape) {
  // Validate ranks.
  ShapeHandle unused_shape;
  TF_RETURN_IF_ERROR(c->WithRank(indices_shape, 2, &unused_shape));
  TF_RETURN_IF_ERROR(c->WithRank(values_shape, 1, &unused_shape));
  TF_RETURN_IF_ERROR(c->WithRank(shape_shape, 1, &unused_shape));

  // Number of elements in indices and values must match.
  DimensionHandle num_index_elements_dim = c->Dim(indices_shape, 0);
  if (c->ValueKnown(num_index_elements_dim)) {
    DimensionHandle num_values_elements_dim = c->Dim(values_shape, 0);
    if (c->ValueKnown(num_values_elements_dim)) {
      int64_t num_index_elements = c->Value(num_index_elements_dim);
      int64_t num_values_elements = c->Value(num_values_elements_dim);
      if (num_index_elements != num_values_elements) {
        return errors::InvalidArgument("Number of elements in index (",
                                       num_index_elements, ") and values (",
                                       num_values_elements, ") do not match.");
      }
    }
  }

  // Rank embedded in indices must match shape.
  DimensionHandle index_rank_dim = c->Dim(indices_shape, 1);
  if (c->ValueKnown(index_rank_dim)) {
    DimensionHandle shape_rank_dim = c->Dim(shape_shape, 0);
    if (c->ValueKnown(shape_rank_dim)) {
      int64_t index_rank = c->Value(index_rank_dim);
      int32_t shape_rank = c->Value(shape_rank_dim);
      if (index_rank != shape_rank) {
        return errors::InvalidArgument("Index rank (", index_rank,
                                       ") and shape rank (", shape_rank,
                                       ") do not match.");
      }
    }
  }

  return Status::OK();
}


Status RandomShape(shape_inference::InferenceContext* c) {
  shape_inference::ShapeHandle out;
  TF_RETURN_IF_ERROR(c->MakeShapeFromShapeTensor(0, &out));
  c->set_output(0, out);
  return Status::OK();
}


Status AvgPoolShape(shape_inference::InferenceContext* c) {
  string data_format_str;
  TensorFormat data_format;
  Status s = c->GetAttr("data_format", &data_format_str);
  if (s.ok()) {
    FormatFromString(data_format_str, &data_format);
  } else {
    data_format = FORMAT_NHWC;
  }

  const int rank = (data_format == FORMAT_NCHW_VECT_C) ? 5 : 4;
  ShapeHandle input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &input_shape));

  TF_RETURN_IF_ERROR(
      CheckFormatConstraintsOnShape(data_format, input_shape, "input", c));

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));
  if (strides.size() != 4) {
    return errors::InvalidArgument(
        "AvgPool requires the stride attribute to contain 4 values, but got: ",
        strides.size());
  }

  std::vector<int32> kernel_sizes;
  TF_RETURN_IF_ERROR(c->GetAttr("ksize", &kernel_sizes));
  if (kernel_sizes.size() != 4) {
    return errors::InvalidArgument(
        "AvgPool requires the ksize attribute to contain 4 values, but got: ",
        kernel_sizes.size());
  }

  int32_t stride_rows = GetTensorDim(strides, data_format, 'H');
  int32_t stride_cols = GetTensorDim(strides, data_format, 'W');
  int32_t kernel_rows = GetTensorDim(kernel_sizes, data_format, 'H');
  int32_t kernel_cols = GetTensorDim(kernel_sizes, data_format, 'W');

  constexpr int num_spatial_dims = 2;
  DimensionHandle batch_size_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'N'));
  DimensionHandle in_rows_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'H'));
  DimensionHandle in_cols_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'W'));
  DimensionHandle depth_dim = c->Dim(
      input_shape, GetTensorDimIndex<num_spatial_dims>(data_format, 'C'));

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  // TODO(mrry,shlens): Raise an error if the stride would cause
  // information in the input to be ignored. This will require a change
  // in the kernel implementation.

  DimensionHandle output_rows, output_cols;
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_rows_dim, kernel_rows, stride_rows, padding, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDims(
      c, in_cols_dim, kernel_cols, stride_cols, padding, &output_cols));

  ShapeHandle output_shape;
  TF_RETURN_IF_ERROR(MakeShapeFromFormat(data_format, batch_size_dim,
                                         {output_rows, output_cols}, depth_dim,
                                         &output_shape, c));
  c->set_output(0, output_shape);
  return Status::OK();
}


Status FusedBatchNormExShape(shape_inference::InferenceContext* c) {
  TF_RETURN_IF_ERROR(FusedBatchNormV3Shape(c));

  string data_format_str;
  TF_RETURN_IF_ERROR(c->GetAttr("data_format", &data_format_str));
  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }
  ShapeHandle x;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 4, &x));

  int channel_dim_index = GetTensorFeatureDimIndex(4, data_format);
  DimensionHandle channel_dim = c->Dim(x, channel_dim_index);

  // This is a cuDNN implementation constraint.
  if (c->ValueKnown(channel_dim) && c->Value(channel_dim) % 4 != 0) {
    return errors::InvalidArgument(
        "_FusedBatchNormEx channel dimension must be divisible by 4.");
  }

  return Status::OK();
}


Status BroadcastBinaryOpOutputShapeFnHelper(InferenceContext* c,
                                            ShapeHandle shape_x,
                                            ShapeHandle shape_y,
                                            bool incompatible_shape_error,
                                            ShapeHandle* out) {
  CHECK_NOTNULL(out);
  if (!c->RankKnown(shape_x) || !c->RankKnown(shape_y)) {
    *out = c->UnknownShape();
    return Status::OK();
  }
  const int32_t rank_x = c->Rank(shape_x);
  const int32_t rank_y = c->Rank(shape_y);
  const int32_t rank_out = std::max(rank_x, rank_y);

  // To compute the broadcast dimensions, we zip together shape_x and shape_y
  // and
  // pad with 1 to make them the same length.
  std::vector<DimensionHandle> dims;
  DimensionHandle dim_one;
  if (rank_x != rank_y) dim_one = c->MakeDim(1);
  for (int i = 0; i < rank_out; ++i) {
    const auto dim_x = i < (rank_out - rank_x)
                           ? dim_one
                           : c->Dim(shape_x, i - (rank_out - rank_x));
    const bool dim_y_is_one = (i < (rank_out - rank_y));
    const auto dim_y =
        dim_y_is_one ? dim_one : c->Dim(shape_y, i - (rank_out - rank_y));
    if (!c->ValueKnown(dim_x) || !c->ValueKnown(dim_y)) {
      // One or both dimensions is unknown.
      //
      // - If either dimension is greater than 1, we assume that the program is
      // correct, and the other dimension will be broadcast to match it.
      // TODO(cwhipkey): For shape inference, if we eliminate the shape checks
      // in C++ op code, we must still assert that the unknown dim is either 1
      // or the same as the known dim.
      // - If either dimension is 1, the other dimension is the output.
      // - If both are unknown then dimension is unknown
      if (c->Value(dim_x) > 1) {
        if (!incompatible_shape_error) {
          *out = c->UnknownShape();
          return Status::OK();
        }
        dims.push_back(dim_x);
      } else if (c->Value(dim_y) > 1) {
        if (!incompatible_shape_error) {
          *out = c->UnknownShape();
          return Status::OK();
        }
        dims.push_back(dim_y);
      } else if (c->Value(dim_x) == 1) {
        dims.push_back(dim_y);
      } else if (c->Value(dim_y) == 1) {
        dims.push_back(dim_x);
      } else if (dim_y.SameHandle(dim_x)) {
        dims.push_back(dim_x);
      } else if (!c->ValueKnown(dim_x) && !c->ValueKnown(dim_y)) {
        dims.push_back(c->UnknownDim());
      } else {
        if (!incompatible_shape_error) {
          *out = c->UnknownShape();
          return Status::OK();
        }
        dims.push_back(c->UnknownDim());
      }
    } else if (c->Value(dim_x) == 1 || c->Value(dim_y) == 1) {
      if (c->Value(dim_x) == 1 && !dim_y_is_one) {
        // We will broadcast dim_x to dim_y.
        dims.push_back(dim_y);
      } else {
        DCHECK_EQ(c->Value(dim_y), 1);
        // We will broadcast dim_y to dim_x.
        dims.push_back(dim_x);
      }
    } else {
      DimensionHandle dim;
      Status s = c->Merge(dim_x, dim_y, &dim);
      if (!s.ok()) {
        if (!incompatible_shape_error) {
          *out = c->MakeShape({});
          return Status::OK();
        }
        return s;
      }
      dims.push_back(dim);
    }
  }

  *out = c->MakeShape(dims);
  return Status::OK();
}


Status Conv2DShapeImpl(shape_inference::InferenceContext* c,
                       bool supports_explicit_padding) {
  string data_format_str, filter_format_str;
  if (!c->GetAttr("data_format", &data_format_str).ok()) {
    data_format_str = "NHWC";
  }
  if (!c->GetAttr("filter_format", &filter_format_str).ok()) {
    filter_format_str = "HWIO";
  }

  TensorFormat data_format;
  if (!FormatFromString(data_format_str, &data_format)) {
    return errors::InvalidArgument("Invalid data format string: ",
                                   data_format_str);
  }
  FilterTensorFormat filter_format;
  if (!FilterFormatFromString(filter_format_str, &filter_format)) {
    return errors::InvalidArgument("Invalid filter format string: ",
                                   filter_format_str);
  }

  constexpr int num_spatial_dims = 2;
  const int rank = GetTensorDimsFromSpatialDims(num_spatial_dims, data_format);
  ShapeHandle conv_input_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(0), rank, &conv_input_shape));
  TF_RETURN_IF_ERROR(CheckFormatConstraintsOnShape(
      data_format, conv_input_shape, "conv_input", c));

  // The filter rank should match the input (4 for NCHW, 5 for NCHW_VECT_C).
  ShapeHandle filter_shape;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), rank, &filter_shape));
  TF_RETURN_IF_ERROR(
      CheckFormatConstraintsOnShape(data_format, filter_shape, "filter", c));

  std::vector<int32> dilations;
  TF_RETURN_IF_ERROR(c->GetAttr("dilations", &dilations));

  if (dilations.size() != 4) {
    return errors::InvalidArgument(
        "Conv2D requires the dilation attribute to contain 4 values, but got: ",
        dilations.size());
  }

  std::vector<int32> strides;
  TF_RETURN_IF_ERROR(c->GetAttr("strides", &strides));

  // strides.size() should be 4 (NCHW) even if the input is 5 (NCHW_VECT_C).
  if (strides.size() != 4) {
    return errors::InvalidArgument("Conv2D on data format ", data_format_str,
                                   " requires the stride attribute to contain"
                                   " 4 values, but got: ",
                                   strides.size());
  }

  const int32_t stride_rows = GetTensorDim(strides, data_format, 'H');
  const int32_t stride_cols = GetTensorDim(strides, data_format, 'W');
  const int32_t dilation_rows = GetTensorDim(dilations, data_format, 'H');
  const int32_t dilation_cols = GetTensorDim(dilations, data_format, 'W');

  DimensionHandle batch_size_dim;
  DimensionHandle input_depth_dim;
  gtl::InlinedVector<DimensionHandle, 2> input_spatial_dims(2);
  TF_RETURN_IF_ERROR(DimensionsFromShape(
      conv_input_shape, data_format, &batch_size_dim,
      absl::MakeSpan(input_spatial_dims), &input_depth_dim, c));

  DimensionHandle output_depth_dim = c->Dim(
      filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'O'));
  DimensionHandle filter_rows_dim = c->Dim(
      filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'H'));
  DimensionHandle filter_cols_dim = c->Dim(
      filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'W'));
  DimensionHandle filter_input_depth_dim;
  if (filter_format == FORMAT_OIHW_VECT_I) {
    TF_RETURN_IF_ERROR(c->Multiply(
        c->Dim(filter_shape,
               GetFilterDimIndex<num_spatial_dims>(filter_format, 'I')),
        c->Dim(filter_shape,
               GetFilterTensorInnerInputChannelsDimIndex(rank, filter_format)),
        &filter_input_depth_dim));
  } else {
    filter_input_depth_dim = c->Dim(
        filter_shape, GetFilterDimIndex<num_spatial_dims>(filter_format, 'I'));
  }

  // Check that the input tensor and the filter tensor agree on the channel
  // count.
  if (c->ValueKnown(input_depth_dim) && c->ValueKnown(filter_input_depth_dim)) {
    int64_t input_depth_value = c->Value(input_depth_dim),
            filter_input_depth_value = c->Value(filter_input_depth_dim);
    if (filter_input_depth_value == 0)
      return errors::InvalidArgument("Depth of filter must not be 0");
    if (input_depth_value % filter_input_depth_value != 0)
      return errors::InvalidArgument(
          "Depth of input (", input_depth_value,
          ") is not a multiple of input depth of filter (",
          filter_input_depth_value, ")");
    if (input_depth_value != filter_input_depth_value) {
      int64_t num_groups = input_depth_value / filter_input_depth_value;
      if (c->ValueKnown(output_depth_dim)) {
        int64_t output_depth_value = c->Value(output_depth_dim);
        if (num_groups == 0)
          return errors::InvalidArgument("Number of groups must not be 0");
        if (output_depth_value % num_groups != 0)
          return errors::InvalidArgument(
              "Depth of output (", output_depth_value,
              ") is not a multiple of the number of groups (", num_groups, ")");
      }
    }
  }

  Padding padding;
  TF_RETURN_IF_ERROR(c->GetAttr("padding", &padding));

  std::vector<int64_t> explicit_paddings;
  if (supports_explicit_padding) {
    Status s = c->GetAttr("explicit_paddings", &explicit_paddings);
    // Use the default value, which is an empty list, if the attribute is not
    // found. Otherwise return the error to the caller.
    if (!s.ok() && !errors::IsNotFound(s)) {
      return s;
    }
    TF_RETURN_IF_ERROR(CheckValidPadding(padding, explicit_paddings,
                                         /*num_dims=*/4, data_format));
  } else {
    CHECK(padding != Padding::EXPLICIT);  // Crash ok.
  }

  DimensionHandle output_rows, output_cols;
  int64_t pad_rows_before = -1, pad_rows_after = -1;
  int64_t pad_cols_before = -1, pad_cols_after = -1;
  if (padding == Padding::EXPLICIT) {
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'H',
                             &pad_rows_before, &pad_rows_after);
    GetExplicitPaddingForDim(explicit_paddings, data_format, 'W',
                             &pad_cols_before, &pad_cols_after);
  }
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, input_spatial_dims[0], filter_rows_dim, dilation_rows, stride_rows,
      padding, pad_rows_before, pad_rows_after, &output_rows));
  TF_RETURN_IF_ERROR(GetWindowedOutputSizeFromDimsV2(
      c, input_spatial_dims[1], filter_cols_dim, dilation_cols, stride_cols,
      padding, pad_cols_before, pad_cols_after, &output_cols));

  absl::optional<DimensionHandle> vect_size;
  if (data_format == FORMAT_NCHW_VECT_C) {
    vect_size.emplace(c->Dim(conv_input_shape,
                             GetTensorInnerFeatureDimIndex(rank, data_format)));
  }
  ShapeHandle output_shape;
  TF_RETURN_IF_ERROR(ShapeFromDimensions(
      batch_size_dim, {output_rows, output_cols}, output_depth_dim, data_format,
      vect_size, c, &output_shape));
  c->set_output(0, output_shape);
  return Status::OK();
}


Status QuantizedAvgPoolShape(InferenceContext* c) {
  TF_RETURN_IF_ERROR(shape_inference::AvgPoolShape(c));
  ShapeHandle unused;
  TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 0, &unused));
  TF_RETURN_IF_ERROR(c->WithRank(c->input(2), 0, &unused));
  c->set_output(1, c->Scalar());
  c->set_output(2, c->Scalar());
  return Status::OK();
}


Status ReductionShape(InferenceContext* c) {
  ShapeHandle input = c->input(0);

  ShapeHandle indices;
  // Older versions of TensorFlow accidentally allowed higher rank tensors like
  // [[1,2]] or [[1],[2]] to represent axis=[1,2].
  if (c->graph_def_version() < 21) {
    indices = c->input(1);
  } else {
    TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &indices));
  }

  bool keep_dims;
  TF_RETURN_IF_ERROR(c->GetAttr("keep_dims", &keep_dims));

  const Tensor* reduction_indices_t = c->input_tensor(1);
  if (reduction_indices_t == nullptr || !c->RankKnown(input)) {
    // If we do not have the reduction values at runtime, or the
    // rank of the input, we don't know the output shape.

    if (keep_dims && c->RankKnown(input)) {
      // output rank matches input input if <keep_dims>.
      c->set_output(0, c->UnknownShapeOfRank(c->Rank(input)));
      return Status::OK();
    } else {
      return shape_inference::UnknownShape(c);
    }
  }

  const int32_t input_rank = c->Rank(input);
  std::set<int64_t> true_indices;
  if (reduction_indices_t->dtype() == DataType::DT_INT32) {
    TF_RETURN_IF_ERROR(ReductionShapeHelper<int32>(reduction_indices_t,
                                                   input_rank, &true_indices));
  } else if (reduction_indices_t->dtype() == DataType::DT_INT64) {
    TF_RETURN_IF_ERROR(ReductionShapeHelper<int64_t>(
        reduction_indices_t, input_rank, &true_indices));
  } else {
    return errors::InvalidArgument(
        "reduction_indices can only be int32 or int64");
  }

  std::vector<DimensionHandle> dims;
  for (int i = 0; i < input_rank; ++i) {
    if (true_indices.count(i) > 0) {
      if (keep_dims) {
        dims.emplace_back(c->MakeDim(1));
      }
    } else {
      dims.emplace_back(c->Dim(input, i));
    }
  }

  c->set_output(0, c->MakeShape(dims));
  return Status::OK();
}


Status BiasAddGradShape(shape_inference::InferenceContext* c) {
  ShapeHandle input_shape;
  // Fetch the data_format attribute, which may not exist.
  string data_format;
  Status s = c->GetAttr("data_format", &data_format);

  if (s.ok() && data_format == "NCHW") {
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 3, &input_shape));
    c->set_output(0, c->Vector(c->Dim(input_shape, 1)));
  } else {
    TF_RETURN_IF_ERROR(c->WithRankAtLeast(c->input(0), 2, &input_shape));
    c->set_output(0, c->Vector(c->Dim(input_shape, -1)));
  }

  return Status::OK();
}


Status MaxPool3DGradShape(shape_inference::InferenceContext* c) {
  return UnchangedShapeWithRank(c, 5);
}


Status MaxPoolShape(shape_inference::InferenceContext* c) {
  return MaxPoolShapeImpl(c, /*supports_explicit_padding=*/false);
}


GF_Err iinf_on_child_box(GF_Box *s, GF_Box *a, Bool is_rem)
{
	GF_ItemInfoBox *ptr = (GF_ItemInfoBox *)s;

	if (a->type == GF_ISOM_BOX_TYPE_INFE) {
		BOX_FIELD_LIST_ASSIGN(item_infos)
		return GF_OK;
	} else {
		return GF_OK;
	}
}

The binary MP4Box in Gpac 1.0.1 has a double-free vulnerability in the gf_list_del function in list.c, which allows attackers to cause a denial of service.


void iloc_box_del(GF_Box *s)
{
	u32 i, item_count;
	GF_ItemLocationBox *ptr = (GF_ItemLocationBox *)s;
	if (ptr == NULL) return;
	item_count = gf_list_count(ptr->location_entries);
	for (i = 0; i < item_count; i++) {
		GF_ItemLocationEntry *location = (GF_ItemLocationEntry *)gf_list_get(ptr->location_entries, i);
		iloc_entry_del(location);
	}
	gf_list_del(ptr->location_entries);
	gf_free(ptr);
}


void iinf_box_del(GF_Box *s)
{
	GF_ItemInfoBox *ptr = (GF_ItemInfoBox *)s;
	if (ptr == NULL) return;
	gf_list_del(ptr->item_infos);
	gf_free(ptr);
}


GF_Err bxml_box_read(GF_Box *s, GF_BitStream *bs)
{
	GF_BinaryXMLBox *ptr = (GF_BinaryXMLBox *)s;

	ptr->data_length = (u32)(ptr->size);
	ptr->data = (char*)gf_malloc(sizeof(char)*ptr->data_length);
	if (!ptr->data) return GF_OUT_OF_MEM;
	gf_bs_read_data(bs, ptr->data, ptr->data_length);
	return GF_OK;
}


GF_Err ireftype_box_write(GF_Box *s, GF_BitStream *bs)
{
	GF_Err e;
	u32 i;
	GF_ItemReferenceTypeBox *ptr = (GF_ItemReferenceTypeBox *)s;
	ptr->type = ptr->reference_type;
	e = gf_isom_box_write_header(s, bs);
	ptr->type = GF_ISOM_BOX_TYPE_REFI;
	if (e) return e;
	gf_bs_write_u16(bs, ptr->from_item_id);
	gf_bs_write_u16(bs, ptr->reference_count);
	for (i = 0; i < ptr->reference_count; i++) {
		gf_bs_write_u16(bs, ptr->to_item_IDs[i]);
	}
	return GF_OK;
}


GF_Err ireftype_box_read(GF_Box *s, GF_BitStream *bs)
{
	u32 i;
	GF_ItemReferenceTypeBox *ptr = (GF_ItemReferenceTypeBox *)s;

	ISOM_DECREASE_SIZE(ptr, 4)
	ptr->from_item_id = gf_bs_read_u16(bs);
	ptr->reference_count = gf_bs_read_u16(bs);
	if (ptr->size / 2 < ptr->reference_count)
		return GF_ISOM_INVALID_FILE;

	ptr->to_item_IDs = (u32 *)gf_malloc(ptr->reference_count * sizeof(u32));
	if (!ptr->to_item_IDs) return GF_OUT_OF_MEM;

	for (i=0; i < ptr->reference_count; i++) {
		ISOM_DECREASE_SIZE(ptr, 2)
		ptr->to_item_IDs[i] = gf_bs_read_u16(bs);
	}
	return GF_OK;
}


GF_Err iinf_box_read(GF_Box *s, GF_BitStream *bs)
{
	GF_ItemInfoBox *ptr = (GF_ItemInfoBox *)s;

	if (ptr->version == 0) {
		ISOM_DECREASE_SIZE(s, 2)
		gf_bs_read_u16(bs);
	} else {
		ISOM_DECREASE_SIZE(s, 4)
		gf_bs_read_u32(bs);
	}
	return gf_isom_box_array_read(s, bs);
}


GF_Err iinf_box_write(GF_Box *s, GF_BitStream *bs)
{
	u32 count;
	GF_Err e;
	GF_ItemInfoBox *ptr = (GF_ItemInfoBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	count = gf_list_count(ptr->item_infos);
	if (ptr->version == 0)
		gf_bs_write_u16(bs, count);
	else
		gf_bs_write_u32(bs, count);

	return GF_OK;
}


GF_Err infe_box_read(GF_Box *s, GF_BitStream *bs)
{
	char *buf;
	u32 buf_len, i, string_len, string_start;
	GF_ItemInfoEntryBox *ptr = (GF_ItemInfoEntryBox *)s;

	ISOM_DECREASE_SIZE(ptr, 4);
	ptr->item_ID = gf_bs_read_u16(bs);
	ptr->item_protection_index = gf_bs_read_u16(bs);

	if (ptr->version == 2) {
		ISOM_DECREASE_SIZE(ptr, 4);
		ptr->item_type = gf_bs_read_u32(bs);
	}
	buf_len = (u32) (ptr->size);
	buf = (char*)gf_malloc(buf_len);
	if (!buf) return GF_OUT_OF_MEM;
	if (buf_len != gf_bs_read_data(bs, buf, buf_len)) {
		gf_free(buf);
		return GF_ISOM_INVALID_FILE;
	}
	string_len = 1;
	string_start = 0;
	for (i = 0; i < buf_len; i++) {
		if (buf[i] == 0) {
			if (!ptr->item_name) {
				ptr->item_name = (char*)gf_malloc(sizeof(char)*string_len);
				if (!ptr->item_name) return GF_OUT_OF_MEM;
				memcpy(ptr->item_name, buf+string_start, string_len);
			} else if (!ptr->content_type) {
				ptr->content_type = (char*)gf_malloc(sizeof(char)*string_len);
				if (!ptr->content_type) return GF_OUT_OF_MEM;
				memcpy(ptr->content_type, buf+string_start, string_len);
			} else if (!ptr->content_encoding) {
				ptr->content_encoding = (char*)gf_malloc(sizeof(char)*string_len);
				if (!ptr->content_encoding) return GF_OUT_OF_MEM;
				memcpy(ptr->content_encoding, buf+string_start, string_len);
			} else {
				//we could throw an error but we silently accept this infe
				break;
			}
			string_start += string_len;
			string_len = 0;
			if (ptr->content_encoding && ptr->version == 1) {
				break;
			}
		}
		string_len++;
	}
	gf_free(buf);
	if (!ptr->item_name || (!ptr->content_type && ptr->version < 2)) {
		GF_LOG(GF_LOG_WARNING, GF_LOG_CONTAINER, ("[isoff] Infe without name or content type !\n"));
	}
	return GF_OK;
}


void ireftype_box_del(GF_Box *s)
{
	GF_ItemReferenceTypeBox *ptr = (GF_ItemReferenceTypeBox *)s;
	if (!ptr) return;
	if (ptr->to_item_IDs) gf_free(ptr->to_item_IDs);
	gf_free(ptr);
}


GF_Err iref_box_write(GF_Box *s, GF_BitStream *bs)
{
	return gf_isom_full_box_write(s, bs);
}


GF_Err infe_box_size(GF_Box *s)
{
	GF_ItemInfoEntryBox *ptr = (GF_ItemInfoEntryBox *)s;
	if (!s) return GF_BAD_PARAM;
	if (ptr->item_type) {
		ptr->version = 2;
		if (ptr->item_ID > 0xFFFF) {
			ptr->version = 3;
		}
	}
	else {
		ptr->version = 0;
	}
	ptr->size += 4;
	if (ptr->version == 3) {
		ptr->size += 2; // item_ID on 32 bits (+2 bytes)
	}
	if (ptr->version >= 2) {
		ptr->size += 4; // item_type size
	}
	if (ptr->item_name) ptr->size += strlen(ptr->item_name)+1;
	else ptr->size += 1;
	if (ptr->item_type == GF_META_ITEM_TYPE_MIME || ptr->item_type == GF_META_ITEM_TYPE_URI) {
		if (ptr->content_type) ptr->size += strlen(ptr->content_type) + 1;
		else ptr->size += 1;
	}
	if (ptr->item_type == GF_META_ITEM_TYPE_MIME) {
		if (ptr->content_encoding) ptr->size += strlen(ptr->content_encoding) + 1;
		else ptr->size += 1;
	}
	return GF_OK;
}


GF_Box *xml_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_XMLBox, GF_ISOM_BOX_TYPE_XML);
	return (GF_Box *)tmp;
}


GF_Err infe_box_write(GF_Box *s, GF_BitStream *bs)
{
	GF_Err e;
	u32 len;
	GF_ItemInfoEntryBox *ptr = (GF_ItemInfoEntryBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	if (ptr->version == 3) {
		gf_bs_write_u32(bs, ptr->item_ID);
	}
	else {
		gf_bs_write_u16(bs, ptr->item_ID);
	}
	gf_bs_write_u16(bs, ptr->item_protection_index);
	if (ptr->version >= 2) {
		gf_bs_write_u32(bs, ptr->item_type);
	}
	if (ptr->item_name) {
		len = (u32) strlen(ptr->item_name)+1;
		gf_bs_write_data(bs, ptr->item_name, len);
	} else {
		gf_bs_write_byte(bs, 0, 1);
	}
	if (ptr->item_type == GF_META_ITEM_TYPE_MIME || ptr->item_type == GF_META_ITEM_TYPE_URI) {
		if (ptr->content_type) {
			len = (u32)strlen(ptr->content_type) + 1;
			gf_bs_write_data(bs, ptr->content_type, len);
		}
		else {
			gf_bs_write_byte(bs, 0, 1);
		}
	}
	if (ptr->item_type == GF_META_ITEM_TYPE_MIME) {
		if (ptr->content_encoding) {
			len = (u32)strlen(ptr->content_encoding) + 1;
			gf_bs_write_data(bs, ptr->content_encoding, len);
		}
		else {
			gf_bs_write_byte(bs, 0, 1);
		}
	}
	return GF_OK;
}


GF_Box *bxml_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_BinaryXMLBox, GF_ISOM_BOX_TYPE_BXML);
	return (GF_Box *)tmp;
}


GF_Err iloc_box_size(GF_Box *s)
{
	u32 i, item_count, extent_count;
	GF_ItemLocationBox *ptr = (GF_ItemLocationBox *)s;
	if (!s) return GF_BAD_PARAM;
	if (ptr->index_size) {
		ptr->version = 1;
	}
	item_count = gf_list_count(ptr->location_entries);
	for (i = 0; i < item_count; i++) {
		GF_ItemLocationEntry *location = (GF_ItemLocationEntry *)gf_list_get(ptr->location_entries, i);
		if (location->construction_method) {
			ptr->version = 1;
		}
		if (location->item_ID > 0xFFFF) {
			ptr->version = 2;
		}
	}
	ptr->size += 4;
	if (ptr->version == 2) {
		ptr->size += 2; // 32 bits item count instead of 16 bits
	}
	for (i = 0; i < item_count; i++) {
		GF_ItemLocationEntry *location = (GF_ItemLocationEntry *)gf_list_get(ptr->location_entries, i);
		extent_count = gf_list_count(location->extent_entries);
		ptr->size += 6 + ptr->base_offset_size + extent_count*(ptr->offset_size + ptr->length_size);
		if (ptr->version == 2) {
			ptr->size += 2; //32 bits item ID instead of 16 bits
		}
		if (ptr->version == 1 || ptr->version == 2) {
			ptr->size += 2; // construction_method
			ptr->size += extent_count*ptr->index_size;
		}
	}
	return GF_OK;
}


void infe_box_del(GF_Box *s)
{
	GF_ItemInfoEntryBox *ptr = (GF_ItemInfoEntryBox *)s;
	if (ptr == NULL) return;
	if (ptr->item_name) gf_free(ptr->item_name);
	if (ptr->full_path) gf_free(ptr->full_path);
	if (ptr->content_type) gf_free(ptr->content_type);
	if (ptr->content_encoding) gf_free(ptr->content_encoding);
	gf_free(ptr);
}


GF_Err xml_box_read(GF_Box *s, GF_BitStream *bs)
{
	GF_XMLBox *ptr = (GF_XMLBox *)s;
	return gf_isom_read_null_terminated_string(s, bs, s->size, &ptr->xml);
}


GF_Box *infe_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_ItemInfoEntryBox, GF_ISOM_BOX_TYPE_INFE);
	return (GF_Box *)tmp;
}


GF_Err iloc_box_read(GF_Box *s, GF_BitStream *bs)
{
	u32 item_count, extent_count, i, j;
	GF_ItemLocationBox *ptr = (GF_ItemLocationBox *)s;

	ISOM_DECREASE_SIZE(ptr, 2)
	ptr->offset_size = gf_bs_read_int(bs, 4);
	ptr->length_size = gf_bs_read_int(bs, 4);
	ptr->base_offset_size = gf_bs_read_int(bs, 4);
	if (ptr->version == 1 || ptr->version == 2) {
		ptr->index_size = gf_bs_read_int(bs, 4);
	} else {
		gf_bs_read_int(bs, 4);
	}
	if (ptr->version < 2) {
		ISOM_DECREASE_SIZE(ptr, 2)
		item_count = gf_bs_read_u16(bs);
	} else {
		ISOM_DECREASE_SIZE(ptr, 4)
		item_count = gf_bs_read_u32(bs);
	}

	for (i = 0; i < item_count; i++) {
		GF_ItemLocationEntry *location_entry;
		GF_SAFEALLOC(location_entry, GF_ItemLocationEntry);
		if (!location_entry) return GF_OUT_OF_MEM;

		gf_list_add(ptr->location_entries, location_entry);
		if (ptr->version < 2) {
			ISOM_DECREASE_SIZE(ptr, 2)
			location_entry->item_ID = gf_bs_read_u16(bs);
		} else {
			ISOM_DECREASE_SIZE(ptr, 4)
			location_entry->item_ID = gf_bs_read_u32(bs);
		}
		if (ptr->version == 1 || ptr->version == 2) {
			ISOM_DECREASE_SIZE(ptr, 2)
			location_entry->construction_method = gf_bs_read_u16(bs);
		}
		else {
			location_entry->construction_method = 0;
		}
		ISOM_DECREASE_SIZE(ptr, (2 + ptr->base_offset_size) )
		location_entry->data_reference_index = gf_bs_read_u16(bs);
		location_entry->base_offset = gf_bs_read_int(bs, 8*ptr->base_offset_size);
#ifndef GPAC_DISABLE_ISOM_WRITE
		location_entry->original_base_offset = location_entry->base_offset;
#endif

		ISOM_DECREASE_SIZE(ptr, 2)
		extent_count = gf_bs_read_u16(bs);
		location_entry->extent_entries = gf_list_new();
		for (j = 0; j < extent_count; j++) {
			GF_ItemExtentEntry *extent_entry;
			GF_SAFEALLOC(extent_entry, GF_ItemExtentEntry);
			if (!extent_entry) return GF_OUT_OF_MEM;

			gf_list_add(location_entry->extent_entries, extent_entry);
			if ((ptr->version == 1 || ptr->version == 2) && ptr->index_size > 0) {
				ISOM_DECREASE_SIZE(ptr, ptr->index_size)
				extent_entry->extent_index = gf_bs_read_int(bs, 8 * ptr->index_size);
			}
			else {
				extent_entry->extent_index = 0;
			}
			ISOM_DECREASE_SIZE(ptr, (ptr->offset_size+ptr->length_size) )

			extent_entry->extent_offset = gf_bs_read_int(bs, 8*ptr->offset_size);
			extent_entry->extent_length = gf_bs_read_int(bs, 8*ptr->length_size);
#ifndef GPAC_DISABLE_ISOM_WRITE
			extent_entry->original_extent_offset = extent_entry->extent_offset;
#endif
		}
	}
	return GF_OK;
}


GF_Err meta_on_child_box(GF_Box *s, GF_Box *a, Bool is_rem)
{
	GF_MetaBox *ptr = (GF_MetaBox *)s;
	switch (a->type) {
	case GF_ISOM_BOX_TYPE_HDLR:
		BOX_FIELD_ASSIGN(handler, GF_HandlerBox)
		break;
	case GF_ISOM_BOX_TYPE_PITM:
		BOX_FIELD_ASSIGN(primary_resource, GF_PrimaryItemBox)
		break;
	case GF_ISOM_BOX_TYPE_DINF:
		BOX_FIELD_ASSIGN(file_locations, GF_DataInformationBox)
		break;
	case GF_ISOM_BOX_TYPE_ILOC:
		BOX_FIELD_ASSIGN(item_locations, GF_ItemLocationBox)
		break;
	case GF_ISOM_BOX_TYPE_IPRO:
		BOX_FIELD_ASSIGN(protections, GF_ItemProtectionBox)
		break;
	case GF_ISOM_BOX_TYPE_IINF:
		BOX_FIELD_ASSIGN(item_infos, GF_ItemInfoBox)
		break;
	case GF_ISOM_BOX_TYPE_IREF:
		BOX_FIELD_ASSIGN(item_refs, GF_ItemReferenceBox);
		break;
	case GF_ISOM_BOX_TYPE_IPRP:
		BOX_FIELD_ASSIGN(item_props, GF_ItemPropertiesBox)
		break;
	case GF_ISOM_BOX_TYPE_GRPL:
		BOX_FIELD_ASSIGN(groups_list, GF_GroupListBox)
		break;
	}
	return GF_OK;
}


GF_Err meta_box_read(GF_Box *s, GF_BitStream *bs)
{
	u64 pos = gf_bs_get_position(bs);
	u64 size = s->size;
	GF_Err e = gf_isom_box_array_read(s, bs);
	/*try to hack around QT files which don't use a full box for meta, rewind 4 bytes*/
	if (e && (pos>4) ) {
		gf_bs_seek(bs, pos-4);
		meta_reset(s);
		s->size = size+4;
		e = gf_isom_box_array_read(s, bs);
	}
	return e;
}


GF_Box *iref_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_ItemReferenceBox, GF_ISOM_BOX_TYPE_IREF);
	tmp->references = gf_list_new();
	return (GF_Box *)tmp;
}


GF_Err bxml_box_size(GF_Box *s)
{
	GF_BinaryXMLBox *ptr = (GF_BinaryXMLBox *)s;
	if (!s) return GF_BAD_PARAM;
	ptr->size += ptr->data_length;
	return GF_OK;
}


void meta_box_del(GF_Box *s)
{
	meta_reset(s);
	gf_free(s);
}


GF_Err iref_box_read(GF_Box *s, GF_BitStream *bs)
{
	return gf_isom_box_array_read_ex(s, bs, s->type);
}


GF_Box *ipro_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_ItemProtectionBox, GF_ISOM_BOX_TYPE_IPRO);
	tmp->protection_information = gf_list_new();
	return (GF_Box *)tmp;
}


GF_Err ipro_on_child_box(GF_Box *s, GF_Box *a, Bool is_rem)
{
	GF_ItemProtectionBox *ptr = (GF_ItemProtectionBox *)s;
	if (a->type == GF_ISOM_BOX_TYPE_SINF) {
		BOX_FIELD_LIST_ASSIGN(protection_information)
		return GF_OK;
	}
	return GF_OK;
}


GF_Err iinf_box_size(GF_Box *s)
{
	u32 pos=0;
	GF_ItemInfoBox *ptr = (GF_ItemInfoBox *)s;
	if (!s) return GF_BAD_PARAM;
	ptr->size += (ptr->version == 0) ? 2 : 4;
	gf_isom_check_position_list(s, ptr->item_infos, &pos);
	return GF_OK;
}


GF_Err xml_box_size(GF_Box *s)
{
	GF_XMLBox *ptr = (GF_XMLBox *)s;
	if (!s) return GF_BAD_PARAM;
    ptr->size += (ptr->xml ? strlen(ptr->xml) : 0) + 1;
	return GF_OK;
}


GF_Box *iloc_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_ItemLocationBox, GF_ISOM_BOX_TYPE_ILOC);
	tmp->location_entries = gf_list_new();
	return (GF_Box *)tmp;
}


GF_Err iloc_box_write(GF_Box *s, GF_BitStream *bs)
{
	GF_Err e;
	u32 i, j, item_count, extent_count;
	GF_ItemLocationBox *ptr = (GF_ItemLocationBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	gf_bs_write_int(bs, ptr->offset_size, 4);
	gf_bs_write_int(bs, ptr->length_size, 4);
	gf_bs_write_int(bs, ptr->base_offset_size, 4);
	gf_bs_write_int(bs, ptr->index_size, 4);
	item_count = gf_list_count(ptr->location_entries);
	if (ptr->version < 2) {
		gf_bs_write_u16(bs, item_count);
	}
	else {
		gf_bs_write_u32(bs, item_count);
	}
	for (i = 0; i < item_count; i++) {
		GF_ItemLocationEntry *location = (GF_ItemLocationEntry *)gf_list_get(ptr->location_entries, i);
		if (ptr->version < 2) {
			gf_bs_write_u16(bs, location->item_ID);
		}
		else {
			gf_bs_write_u32(bs, location->item_ID);
		}
		if (ptr->version == 1 || ptr->version == 2) {
			gf_bs_write_u16(bs, location->construction_method);
		}
		gf_bs_write_u16(bs, location->data_reference_index);
		gf_bs_write_long_int(bs, location->base_offset, 8*ptr->base_offset_size);
		extent_count = gf_list_count(location->extent_entries);
		gf_bs_write_u16(bs, extent_count);
		for (j=0; j<extent_count; j++) {
			GF_ItemExtentEntry *extent = (GF_ItemExtentEntry *)gf_list_get(location->extent_entries, j);
			if ((ptr->version == 1 || ptr->version == 2) && ptr->index_size > 0) {
				gf_bs_write_long_int(bs, extent->extent_index, 8 * ptr->index_size);
			}
			gf_bs_write_long_int(bs, extent->extent_offset, 8*ptr->offset_size);
			gf_bs_write_long_int(bs, extent->extent_length, 8*ptr->length_size);
		}
	}
	return GF_OK;
}


void ipro_box_del(GF_Box *s)
{
	GF_ItemProtectionBox *ptr = (GF_ItemProtectionBox *)s;
	if (ptr == NULL) return;
	gf_list_del(ptr->protection_information);
	gf_free(ptr);
}


GF_Box *ireftype_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_ItemReferenceTypeBox, GF_ISOM_BOX_TYPE_REFI);
	return (GF_Box *)tmp;
}


GF_Err pitm_box_size(GF_Box *s)
{
	GF_ItemLocationBox *ptr = (GF_ItemLocationBox *)s;
	if (!s) return GF_BAD_PARAM;
	ptr->size += 2;
	return GF_OK;
}


GF_Err meta_box_write(GF_Box *s, GF_BitStream *bs)
{
	return gf_isom_full_box_write(s, bs);
}


GF_Box *pitm_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_PrimaryItemBox, GF_ISOM_BOX_TYPE_PITM);
	return (GF_Box *)tmp;
}


GF_Err meta_box_size(GF_Box *s)
{
	u32 pos=0;
	GF_MetaBox *ptr = (GF_MetaBox *)s;

	gf_isom_check_position(s, (GF_Box *)ptr->handler, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->primary_resource, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->file_locations, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->item_locations, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->protections, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->item_infos, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->item_refs, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->item_props, &pos);
	gf_isom_check_position(s, (GF_Box *)ptr->groups_list, &pos);
	return GF_OK;
}


void pitm_box_del(GF_Box *s)
{
	GF_PrimaryItemBox *ptr = (GF_PrimaryItemBox *)s;
	if (ptr == NULL) return;
	gf_free(ptr);
}


GF_Err xml_box_write(GF_Box *s, GF_BitStream *bs)
{
	GF_Err e;
	GF_XMLBox *ptr = (GF_XMLBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	if (ptr->xml) gf_bs_write_data(bs, ptr->xml, (u32) strlen(ptr->xml));
	gf_bs_write_u8(bs, 0);
	return GF_OK;
}


GF_Box *iinf_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_ItemInfoBox, GF_ISOM_BOX_TYPE_IINF);
	tmp->item_infos = gf_list_new();
	return (GF_Box *)tmp;
}


void iloc_entry_del(GF_ItemLocationEntry *location)
{
	u32 j, extent_count;
	extent_count = gf_list_count(location->extent_entries);
	for (j = 0; j < extent_count; j++) {
		GF_ItemExtentEntry *extent = (GF_ItemExtentEntry *)gf_list_get(location->extent_entries, j);
		gf_free(extent);
	}
	gf_list_del(location->extent_entries);
	gf_free(location);
}


GF_Err bxml_box_write(GF_Box *s, GF_BitStream *bs)
{
	GF_Err e;
	GF_BinaryXMLBox *ptr = (GF_BinaryXMLBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	if (ptr->data_length) gf_bs_write_data(bs, ptr->data, ptr->data_length);
	return GF_OK;
}


void meta_reset(GF_Box *s)
{
	GF_MetaBox *ptr = (GF_MetaBox *)s;
	if (ptr == NULL) return;

	//destroy all boxes
	if (ptr->child_boxes) gf_isom_box_array_del(ptr->child_boxes);
	ptr->child_boxes = NULL;

	//reset pointers (destroyed above)
	ptr->handler = NULL;
	ptr->primary_resource = NULL;
	ptr->file_locations = NULL;
	ptr->item_locations = NULL;
	ptr->protections = NULL;
	ptr->item_infos = NULL;
	ptr->IPMP_control = NULL;
	ptr->item_refs = NULL;
	ptr->item_props = NULL;
	ptr->groups_list = NULL;
}


GF_Box *meta_box_new()
{
	ISOM_DECL_BOX_ALLOC(GF_MetaBox, GF_ISOM_BOX_TYPE_META);
	tmp->child_boxes = gf_list_new();
	return (GF_Box *)tmp;
}


GF_Err ipro_box_size(GF_Box *s)
{
	GF_ItemProtectionBox *ptr = (GF_ItemProtectionBox *)s;
	u32 pos=0;
	gf_isom_check_position_list(s, ptr->protection_information, &pos);
	s->size += 2;
	return GF_OK;
}


GF_Err ipro_box_read(GF_Box *s, GF_BitStream *bs)
{
	ISOM_DECREASE_SIZE(s, 2)
	gf_bs_read_u16(bs);
	return gf_isom_box_array_read(s, bs);
}


GF_Err ipro_box_write(GF_Box *s, GF_BitStream *bs)
{
	u32 count;
	GF_Err e;
	GF_ItemProtectionBox *ptr = (GF_ItemProtectionBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	count = gf_list_count(ptr->protection_information);
	gf_bs_write_u16(bs, count);
	return GF_OK;
}


GF_Err pitm_box_read(GF_Box *s, GF_BitStream *bs)
{
	GF_PrimaryItemBox *ptr = (GF_PrimaryItemBox *)s;
	ISOM_DECREASE_SIZE(ptr, 2)
	ptr->item_ID = gf_bs_read_u16(bs);
	return GF_OK;
}


GF_Err iref_on_child_box(GF_Box *s, GF_Box *a, Bool is_rem)
{
	GF_ItemReferenceBox *ptr = (GF_ItemReferenceBox *)s;
	BOX_FIELD_LIST_ASSIGN(references)
	return GF_OK;
}


void xml_box_del(GF_Box *s)
{
	GF_XMLBox *ptr = (GF_XMLBox *)s;
	if (ptr == NULL) return;
	if (ptr->xml) gf_free(ptr->xml);
	gf_free(ptr);
}


GF_Err iref_box_size(GF_Box *s)
{
	u32 pos=0;
	GF_ItemReferenceBox *ptr = (GF_ItemReferenceBox *)s;
	gf_isom_check_position_list(s, ptr->references, &pos);
	return GF_OK;
}


GF_Err ireftype_box_size(GF_Box *s)
{
	GF_ItemReferenceTypeBox *ptr = (GF_ItemReferenceTypeBox *)s;
	ptr->size += 4 + (ptr->reference_count * sizeof(u16));
	return GF_OK;
}


void iref_box_del(GF_Box *s)
{
	GF_ItemReferenceBox *ptr = (GF_ItemReferenceBox *)s;
	if (ptr == NULL) return;
	gf_list_del(ptr->references);
	gf_free(ptr);
}


GF_Err pitm_box_write(GF_Box *s, GF_BitStream *bs)
{
	GF_Err e;
	GF_PrimaryItemBox *ptr = (GF_PrimaryItemBox *)s;
	if (!s) return GF_BAD_PARAM;
	e = gf_isom_full_box_write(s, bs);
	if (e) return e;
	gf_bs_write_u16(bs, ptr->item_ID);
	return GF_OK;
}


void bxml_box_del(GF_Box *s)
{
	GF_BinaryXMLBox *ptr = (GF_BinaryXMLBox *)s;
	if (ptr == NULL) return;
	if (ptr->data) gf_free(ptr->data);
	gf_free(ptr);
}


Prevent overflow in sparse dense cwise ops.

PiperOrigin-RevId: 415543171
Change-Id: I22dab7c41be2121ab5efe5403ca0e2f9b7cb24b8

  void Compute(OpKernelContext *ctx) override {
    const Tensor *indices_t, *values_t, *shape_t, *dense_t;
    OP_REQUIRES_OK(ctx, ctx->input("sp_indices", &indices_t));
    OP_REQUIRES_OK(ctx, ctx->input("sp_values", &values_t));
    OP_REQUIRES_OK(ctx, ctx->input("sp_shape", &shape_t));
    OP_REQUIRES_OK(ctx, ctx->input("dense", &dense_t));

    // Validations.
    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrix(indices_t->shape()),
                errors::InvalidArgument(
                    "Input sp_indices should be a matrix but received shape: ",
                    indices_t->shape().DebugString()));
    OP_REQUIRES(ctx,
                TensorShapeUtils::IsVector(values_t->shape()) &&
                    TensorShapeUtils::IsVector(shape_t->shape()),
                errors::InvalidArgument(
                    "Inputs sp_values and sp_shape should be vectors "
                    "but received shapes: ",
                    values_t->shape().DebugString(), " and ",
                    shape_t->shape().DebugString()));
    OP_REQUIRES(
        ctx, TensorShapeUtils::IsVector(shape_t->shape()),
        errors::InvalidArgument("Input sp_shape must be a vector. Got: ",
                                shape_t->shape().DebugString()));
    OP_REQUIRES(
        ctx, values_t->dim_size(0) == indices_t->dim_size(0),
        errors::InvalidArgument(
            "The first dimension of values and indices should match. (",
            values_t->dim_size(0), " vs. ", indices_t->dim_size(0), ")"));
    OP_REQUIRES(
        ctx, shape_t->shape().dim_size(0) == indices_t->shape().dim_size(1),
        errors::InvalidArgument(
            "Number of dimensions must match second dimension of indices. ",
            "Got ", shape_t->shape().dim_size(0),
            " dimensions, indices shape: ", indices_t->shape().DebugString()));
    OP_REQUIRES(ctx, shape_t->NumElements() > 0,
                errors::InvalidArgument(
                    "The shape argument requires at least one element."));

    const auto indices_mat = indices_t->matrix<int64_t>();
    const auto shape_vec = shape_t->vec<int64_t>();
    TensorShape lhs_shape;
    OP_REQUIRES_OK(ctx, TensorShape::BuildTensorShape(shape_vec, &lhs_shape));
    const auto lhs_dims = BCast::FromShape(lhs_shape);
    const auto rhs_dims = BCast::FromShape(dense_t->shape());
    BCast b(lhs_dims, rhs_dims, false);  // false for keeping the same num dims.

    // True iff (size(lhs) >= size(rhs)) and all dims in lhs is greater or equal
    // to dims in rhs (from right to left).
    auto VecGreaterEq = [](ArraySlice<int64_t> lhs, ArraySlice<int64_t> rhs) {
      if (lhs.size() < rhs.size()) return false;
      for (size_t i = 0; i < rhs.size(); ++i) {
        if (lhs[lhs.size() - 1 - i] < rhs[rhs.size() - 1 - i]) return false;
      }
      return true;
    };
    OP_REQUIRES(ctx, VecGreaterEq(lhs_dims, rhs_dims) && b.IsValid(),
                errors::InvalidArgument(
                    "SparseDenseBinaryOpShared broadcasts dense to sparse "
                    "only; got incompatible shapes: [",
                    absl::StrJoin(lhs_dims, ","), "] vs. [",
                    absl::StrJoin(rhs_dims, ","), "]"));

    Tensor *output_values = nullptr;
    Tensor dense_gathered;
    const int64_t nnz = indices_t->dim_size(0);
    OP_REQUIRES_OK(ctx,
                   ctx->allocate_output(0, TensorShape({nnz}), &output_values));
    OP_REQUIRES_OK(
        ctx, ctx->allocate_temp(DataTypeToEnum<T>::value, TensorShape({nnz}),
                                &dense_gathered));
    bool op_is_div = false;
    if (absl::StrContains(ctx->op_kernel().type_string_view(), "Div")) {
      op_is_div = true;
    }
    // Pulls relevant entries from the dense side, with reshape and broadcasting
    // *of the dense side* taken into account.  Use a TensorRef to avoid blowing
    // up memory.
    //
    // We can directly use the sparse indices to look up dense side, because
    // "b.y_reshape()" and "b.y_bcast()" are guaranteed to have rank "ndims".
    auto dense_gathered_flat = dense_gathered.flat<T>();
    const int ndims = lhs_dims.size();
    switch (ndims) {
#define CASE(NDIM)                                                             \
  case NDIM: {                                                                 \
    TensorRef<Eigen::Tensor<const T, NDIM, Eigen::RowMajor>> rhs_ref =         \
        dense_t->shaped<T, NDIM>(b.y_reshape())                                \
            .broadcast(BCast::ToIndexArray<NDIM>(b.y_bcast()));                \
    Eigen::array<Eigen::DenseIndex, NDIM> idx;                                 \
    bool indices_valid = true;                                                 \
    for (int i = 0; i < nnz; ++i) {                                            \
      for (int d = 0; d < NDIM; ++d) {                                         \
        idx[d] = internal::SubtleMustCopy(indices_mat(i, d));                  \
        if (!FastBoundsCheck(idx[d], rhs_ref.dimension(d))) {                  \
          indices_valid = false;                                               \
        }                                                                      \
      }                                                                        \
      OP_REQUIRES(                                                             \
          ctx, indices_valid,                                                  \
          errors::InvalidArgument("Provided indices are out-of-bounds w.r.t. " \
                                  "dense side with broadcasted shape"));       \
      dense_gathered_flat(i) = rhs_ref.coeff(idx);                             \
      if (op_is_div) {                                                         \
        OP_REQUIRES(ctx, dense_gathered_flat(i) != 0,                          \
                    errors::InvalidArgument(                                   \
                        "SparseDenseCwiseDiv cannot divide by zero,"           \
                        "but input dense tensor contains zero "));             \
      }                                                                        \
    }                                                                          \
    break;                                                                     \
  }

      CASE(1);
      CASE(2);
      CASE(3);
      CASE(4);
      CASE(5);
      default:
        OP_REQUIRES(
            ctx, false,
            errors::InvalidArgument("Only tensors with ranks between 1 and 5 "
                                    "are currently supported.  Tensor rank: ",
                                    ndims));
#undef CASE
    }

    output_values->flat<T>().device(ctx->eigen_device<Device>()) =
        values_t->flat<T>().binaryExpr(dense_gathered_flat,
                                       typename Functor::func());
  }

Tensorflow is an Open Source Machine Learning Framework. The implementations of `Sparse*Cwise*` ops are vulnerable to integer overflows. These can be used to trigger large allocations (so, OOM based denial of service) or `CHECK`-fails when building new `TensorShape` objects (so, assert failures based denial of service). We are missing some validation on the shapes of the input tensors as well as directly constructing a large `TensorShape` with user-provided dimensions. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.


Prevent integer overflow in `OpLevelCostEstimator::CalculateOutputSize`.

In order to not change the API, we return a negative value in case of overflow. A better fix is to change the API to return a status instead.

PiperOrigin-RevId: 408701427
Change-Id: Idf31e7f0bf18ca824d084fdd355e1f653f145c20

int64_t OpLevelCostEstimator::CalculateTensorElementCount(
    const OpInfo::TensorProperties& tensor, bool* found_unknown_shapes) {
  VLOG(2) << "   with " << DataTypeString(tensor.dtype()) << " tensor of shape "
          << tensor.shape().DebugString();
  int64_t tensor_size = 1;
  int num_dims = std::max(1, tensor.shape().dim_size());
  auto tensor_shape =
      MaybeGetMinimumShape(tensor.shape(), num_dims, found_unknown_shapes);
  for (const auto& dim : tensor_shape.dim()) {
    tensor_size *= dim.size();
  }
  return tensor_size;
}

Tensorflow is an Open Source Machine Learning Framework. The implementation of `OpLevelCostEstimator::CalculateOutputSize` is vulnerable to an integer overflow if an attacker can create an operation which would involve tensors with large enough number of elements. We can have a large enough number of dimensions in `output_shape.dim()` or just a small number of dimensions being large enough to cause an overflow in the multiplication. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.


Add negative bound check for row and column pooling_sequence in FractionalAvgPoolGrad op to avoid out of bound heap access

PiperOrigin-RevId: 413837346
Change-Id: I2b86034101df31bee161abcb781755e236c7bccd

  void Compute(OpKernelContext* context) override {
    typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
        ConstEigenMatrixMap;
    typedef Eigen::Map<Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
        EigenMatrixMap;

    constexpr int tensor_in_and_out_dims = 4;

    const Tensor& tensor_in = context->input(0);
    OP_REQUIRES(context, tensor_in.dims() == tensor_in_and_out_dims,
                errors::InvalidArgument("tensor_in must be 4-dimensional"));

    std::vector<int> input_size(tensor_in_and_out_dims);
    std::vector<int> output_size(tensor_in_and_out_dims);
    for (int i = 0; i < tensor_in_and_out_dims; ++i) {
      input_size[i] = tensor_in.dim_size(i);
      OP_REQUIRES(
          context, pooling_ratio_[i] <= input_size[i],
          errors::InvalidArgument(
              "Pooling ratio cannot be bigger than input tensor dim size."));
    }
    // Output size.
    for (int i = 0; i < tensor_in_and_out_dims; ++i) {
      output_size[i] =
          static_cast<int>(std::floor(input_size[i] / pooling_ratio_[i]));
      DCHECK_GT(output_size[i], 0);
    }

    // Generate pooling sequence.
    std::vector<int64_t> row_cum_seq;
    std::vector<int64_t> col_cum_seq;
    GuardedPhiloxRandom generator;
    generator.Init(seed_, seed2_);
    row_cum_seq = GeneratePoolingSequence(input_size[1], output_size[1],
                                          &generator, pseudo_random_);
    col_cum_seq = GeneratePoolingSequence(input_size[2], output_size[2],
                                          &generator, pseudo_random_);

    // Prepare output.
    Tensor* output_tensor = nullptr;
    OP_REQUIRES_OK(context, context->allocate_output(
                                0,
                                TensorShape({output_size[0], output_size[1],
                                             output_size[2], output_size[3]}),
                                &output_tensor));
    Tensor* output_row_seq_tensor = nullptr;
    OP_REQUIRES_OK(
        context, context->allocate_output(
                     1, TensorShape({static_cast<int64_t>(row_cum_seq.size())}),
                     &output_row_seq_tensor));
    Tensor* output_col_seq_tensor = nullptr;
    OP_REQUIRES_OK(
        context, context->allocate_output(
                     2, TensorShape({static_cast<int64_t>(col_cum_seq.size())}),
                     &output_col_seq_tensor));

    ConstEigenMatrixMap in_mat(tensor_in.flat<T>().data(), input_size[3],
                               input_size[2] * input_size[1] * input_size[0]);

    EigenMatrixMap out_mat(output_tensor->flat<T>().data(), output_size[3],
                           output_size[2] * output_size[1] * output_size[0]);
    // out_count corresponds to number of elements in each pooling cell.
    Eigen::Matrix<T, Eigen::Dynamic, 1> out_count(out_mat.cols());

    // Initializes the output tensor and out_count with 0.
    out_mat.setZero();
    out_count.setZero();

    auto output_row_seq_flat = output_row_seq_tensor->flat<int64_t>();
    auto output_col_seq_flat = output_col_seq_tensor->flat<int64_t>();

    // Set output tensors.
    for (int i = 0; i < row_cum_seq.size(); ++i) {
      output_row_seq_flat(i) = row_cum_seq[i];
    }

    for (int i = 0; i < col_cum_seq.size(); ++i) {
      output_col_seq_flat(i) = col_cum_seq[i];
    }

    // For both input and output,
    // 0: batch
    // 1: row / row
    // 2: col / col
    // 3: depth / channel
    const int64_t row_max = input_size[1] - 1;
    const int64_t col_max = input_size[2] - 1;
    for (int64_t b = 0; b < input_size[0]; ++b) {
      // row sequence.
      for (int64_t hs = 0; hs < row_cum_seq.size() - 1; ++hs) {
        // row start and end.
        const int64_t row_start = row_cum_seq[hs];
        int64_t row_end =
            overlapping_ ? row_cum_seq[hs + 1] : row_cum_seq[hs + 1] - 1;
        row_end = std::min(row_end, row_max);

        // col sequence.
        for (int64_t ws = 0; ws < col_cum_seq.size() - 1; ++ws) {
          const int64_t out_offset =
              (b * output_size[1] + hs) * output_size[2] + ws;
          // col start and end.
          const int64_t col_start = col_cum_seq[ws];
          int64_t col_end =
              overlapping_ ? col_cum_seq[ws + 1] : col_cum_seq[ws + 1] - 1;
          col_end = std::min(col_end, col_max);
          for (int64_t h = row_start; h <= row_end; ++h) {
            for (int64_t w = col_start; w <= col_end; ++w) {
              const int64_t in_offset =
                  (b * input_size[1] + h) * input_size[2] + w;
              out_mat.col(out_offset) += in_mat.col(in_offset);
              out_count(out_offset)++;
            }
          }
        }
      }
    }
    DCHECK_GT(out_count.minCoeff(), 0);
    out_mat.array().rowwise() /= out_count.transpose().array();
  }

Tensorflow is an Open Source Machine Learning Framework. The implementation of `FractionalAvgPoolGrad` does not consider cases where the input tensors are invalid allowing an attacker to read from outside of bounds of heap. The fix will be included in TensorFlow 2.8.0. We will also cherrypick this commit on TensorFlow 2.7.1, TensorFlow 2.6.3, and TensorFlow 2.5.3, as these are also affected and still in supported range.


  explicit FractionalAvgPoolGradOp(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("overlapping", &overlapping_));
  }


  explicit FractionalAvgPoolOp(OpKernelConstruction* context)
      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("pooling_ratio", &pooling_ratio_));
    OP_REQUIRES_OK(context, context->GetAttr("pseudo_random", &pseudo_random_));
    OP_REQUIRES_OK(context, context->GetAttr("overlapping", &overlapping_));
    OP_REQUIRES(context, pooling_ratio_.size() == 4,
                errors::InvalidArgument(
                    "pooling_ratio field must specify 4 dimensions"));
    OP_REQUIRES(
        context, pooling_ratio_[0] == 1 || pooling_ratio_[3] == 1,
        errors::Unimplemented("Fractional average pooling is not yet "
                              "supported on the batch nor channel dimension."));
    OP_REQUIRES_OK(context, context->GetAttr("deterministic", &deterministic_));
    OP_REQUIRES_OK(context, context->GetAttr("seed", &seed_));
    OP_REQUIRES_OK(context, context->GetAttr("seed2", &seed2_));
    if (deterministic_) {
      // If both seeds are not set when deterministic_ is true, force set seeds.
      if ((seed_ == 0) && (seed2_ == 0)) {
        seed_ = random::New64();
        seed2_ = random::New64();
      }
    } else {
      OP_REQUIRES(
          context, (seed_ == 0) && (seed2_ == 0),
          errors::InvalidArgument(
              "Both seed and seed2 should be 0 if deterministic is false."));
    }
  }


  void Compute(OpKernelContext* context) override {
    // Here's the basic idea:
    // Batch and depth dimension are independent from row and col dimension. And
    // because FractionalAvgPool currently only support pooling along row and
    // col, we can basically think of this 4D tensor backpropagation as
    // operation of a series of 2D planes.
    //
    // For each element of a 'slice' (2D plane) of output_backprop, we need to
    // figure out its contributors when doing FractionalAvgPool operation. This
    // can be done based on row_pooling_sequence, col_pooling_seq and
    // overlapping.
    // Once we figure out the original contributors, we just need to evenly
    // divide the value of this element among these contributors.
    //
    // Internally, we divide the out_backprop tensor and store it in a temporary
    // tensor of double type. And cast it to the corresponding type.
    typedef Eigen::Map<const Eigen::Matrix<T, Eigen::Dynamic, Eigen::Dynamic>>
        ConstEigenMatrixMap;
    typedef Eigen::Map<Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic>>
        EigenDoubleMatrixMap;

    // Grab the inputs.
    const Tensor& orig_input_tensor_shape = context->input(0);
    OP_REQUIRES(context,
                orig_input_tensor_shape.dims() == 1 &&
                    orig_input_tensor_shape.NumElements() == 4,
                errors::InvalidArgument("original input tensor shape must be"
                                        "1-dimensional and 4 elements"));
    const Tensor& out_backprop = context->input(1);
    const Tensor& row_seq_tensor = context->input(2);
    const Tensor& col_seq_tensor = context->input(3);

    const int64_t out_batch = out_backprop.dim_size(0);
    const int64_t out_rows = out_backprop.dim_size(1);
    const int64_t out_cols = out_backprop.dim_size(2);
    const int64_t out_depth = out_backprop.dim_size(3);

    OP_REQUIRES(context, row_seq_tensor.NumElements() > out_rows,
                errors::InvalidArgument("Given out_backprop shape ",
                                        out_backprop.shape().DebugString(),
                                        ", row_seq_tensor must have at least ",
                                        out_rows + 1, " elements, but got ",
                                        row_seq_tensor.NumElements()));
    OP_REQUIRES(context, col_seq_tensor.NumElements() > out_cols,
                errors::InvalidArgument("Given out_backprop shape ",
                                        out_backprop.shape().DebugString(),
                                        ", col_seq_tensor must have at least ",
                                        out_cols + 1, " elements, but got ",
                                        col_seq_tensor.NumElements()));

    auto row_seq_tensor_flat = row_seq_tensor.flat<int64_t>();
    auto col_seq_tensor_flat = col_seq_tensor.flat<int64_t>();
    auto orig_input_tensor_shape_flat = orig_input_tensor_shape.flat<int64_t>();

    const int64_t in_batch = orig_input_tensor_shape_flat(0);
    const int64_t in_rows = orig_input_tensor_shape_flat(1);
    const int64_t in_cols = orig_input_tensor_shape_flat(2);
    const int64_t in_depth = orig_input_tensor_shape_flat(3);
    OP_REQUIRES(
        context, in_batch != 0,
        errors::InvalidArgument("Batch dimension of input must not be 0"));
    OP_REQUIRES(
        context, in_rows != 0,
        errors::InvalidArgument("Rows dimension of input must not be 0"));
    OP_REQUIRES(
        context, in_cols != 0,
        errors::InvalidArgument("Columns dimension of input must not be 0"));
    OP_REQUIRES(
        context, in_depth != 0,
        errors::InvalidArgument("Depth dimension of input must not be 0"));

    constexpr int tensor_in_and_out_dims = 4;
    // Transform orig_input_tensor_shape into TensorShape
    TensorShape in_shape;
    for (auto i = 0; i < tensor_in_and_out_dims; ++i) {
      in_shape.AddDim(orig_input_tensor_shape_flat(i));
    }

    // Create intermediate in_backprop.
    Tensor in_backprop_tensor_temp;
    OP_REQUIRES_OK(context, context->forward_input_or_allocate_temp(
                                {0}, DataTypeToEnum<double>::v(), in_shape,
                                &in_backprop_tensor_temp));
    in_backprop_tensor_temp.flat<double>().setZero();
    // Transform 4D tensor to 2D matrix.
    EigenDoubleMatrixMap in_backprop_tensor_temp_mat(
        in_backprop_tensor_temp.flat<double>().data(), in_depth,
        in_cols * in_rows * in_batch);
    ConstEigenMatrixMap out_backprop_mat(out_backprop.flat<T>().data(),
                                         out_depth,
                                         out_cols * out_rows * out_batch);
    // Loop through each element of out_backprop and evenly distribute the
    // element to the corresponding pooling cell.
    const int64_t in_max_row_index = in_rows - 1;
    const int64_t in_max_col_index = in_cols - 1;
    for (int64_t b = 0; b < out_batch; ++b) {
      for (int64_t r = 0; r < out_rows; ++r) {
        const int64_t in_row_start = row_seq_tensor_flat(r);

        int64_t in_row_end = overlapping_ ? row_seq_tensor_flat(r + 1)
                                          : row_seq_tensor_flat(r + 1) - 1;
        in_row_end = std::min(in_row_end, in_max_row_index);
        OP_REQUIRES(context, in_row_start >= 0 && in_row_end >= 0,
                    errors::InvalidArgument(
                        "Row sequence tensor values must not be negative, got ",
                        row_seq_tensor_flat));

        for (int64_t c = 0; c < out_cols; ++c) {
          const int64_t in_col_start = col_seq_tensor_flat(c);
          int64_t in_col_end = overlapping_ ? col_seq_tensor_flat(c + 1)
                                            : col_seq_tensor_flat(c + 1) - 1;
          in_col_end = std::min(in_col_end, in_max_col_index);

          OP_REQUIRES(
              context, in_col_start >= 0 && in_col_end >= 0,
              errors::InvalidArgument(
                  "Column sequence tensor values must not be negative, got ",
                  col_seq_tensor_flat));
          const int64_t num_elements_in_pooling_cell =
              (in_row_end - in_row_start + 1) * (in_col_end - in_col_start + 1);
          const int64_t out_index = (b * out_rows + r) * out_cols + c;
          // Now we can evenly distribute out_backprop(b, h, w, *) to
          // in_backprop(b, hs:he, ws:we, *).
          for (int64_t in_r = in_row_start; in_r <= in_row_end; ++in_r) {
            for (int64_t in_c = in_col_start; in_c <= in_col_end; ++in_c) {
              const int64_t in_index = (b * in_rows + in_r) * in_cols + in_c;
              // Walk through each channel (depth).
              for (int64_t d = 0; d < out_depth; ++d) {
                const double out_backprop_element = static_cast<double>(
                    out_backprop_mat.coeffRef(d, out_index));
                double& in_backprop_ref =
                    in_backprop_tensor_temp_mat.coeffRef(d, in_index);
                in_backprop_ref +=
                    out_backprop_element / num_elements_in_pooling_cell;
              }
            }
          }
        }
      }
    }

    // Depending on the type, cast double to type T.
    Tensor* in_backprop_tensor = nullptr;
    OP_REQUIRES_OK(context, context->forward_input_or_allocate_output(
                                {0}, 0, in_shape, &in_backprop_tensor));
    auto in_backprop_tensor_flat = in_backprop_tensor->flat<T>();
    auto in_backprop_tensor_temp_flat = in_backprop_tensor_temp.flat<double>();
    for (int64_t i = 0; i < in_backprop_tensor_flat.size(); ++i) {
      in_backprop_tensor_flat(i) =
          static_cast<T>(in_backprop_tensor_temp_flat(i));
    }
  }


Ensure that bundles have metadata on install

If we have a bundle without metadata we wouldn't properly present
the permissions in the transaction.

flatpak_dir_ensure_bundle_remote (FlatpakDir         *self,
                                  GFile              *file,
                                  GBytes             *extra_gpg_data,
                                  FlatpakDecomposed **out_ref,
                                  char              **out_checksum,
                                  char              **out_metadata,
                                  gboolean           *out_created_remote,
                                  GCancellable       *cancellable,
                                  GError            **error)
{
  g_autoptr(FlatpakDecomposed) ref = NULL;
  gboolean created_remote = FALSE;
  g_autoptr(GBytes) deploy_data = NULL;
  g_autoptr(GVariant) metadata = NULL;
  g_autofree char *origin = NULL;
  g_autofree char *fp_metadata = NULL;
  g_autofree char *basename = NULL;
  g_autoptr(GBytes) included_gpg_data = NULL;
  GBytes *gpg_data = NULL;
  g_autofree char *to_checksum = NULL;
  g_autofree char *remote = NULL;
  g_autofree char *collection_id = NULL;

  if (!flatpak_dir_ensure_repo (self, cancellable, error))
    return NULL;

  metadata = flatpak_bundle_load (file, &to_checksum,
                                  &ref,
                                  &origin,
                                  NULL, &fp_metadata, NULL,
                                  &included_gpg_data,
                                  &collection_id,
                                  error);
  if (metadata == NULL)
    return NULL;

  /* If we rely on metadata (to e.g. print permissions), check it exists before creating the remote */
  if (out_metadata && fp_metadata == NULL)
    {
      flatpak_fail_error (error, FLATPAK_ERROR_INVALID_DATA, "No metadata in bundler header");
      return NULL;
    }

  gpg_data = extra_gpg_data ? extra_gpg_data : included_gpg_data;

  deploy_data = flatpak_dir_get_deploy_data (self, ref, FLATPAK_DEPLOY_VERSION_ANY, cancellable, NULL);
  if (deploy_data != NULL)
    {
      remote = g_strdup (flatpak_deploy_data_get_origin (deploy_data));

      /* We need to import any gpg keys because otherwise the pull will fail */
      if (gpg_data != NULL)
        {
          g_autoptr(GKeyFile) new_config = NULL;

          new_config = ostree_repo_copy_config (flatpak_dir_get_repo (self));

          if (!flatpak_dir_modify_remote (self, remote, new_config,
                                          gpg_data, cancellable, error))
            return NULL;
        }
    }
  else
    {
      g_autofree char *id = flatpak_decomposed_dup_id (ref);
      /* Add a remote for later updates */
      basename = g_file_get_basename (file);
      remote = flatpak_dir_create_origin_remote (self,
                                                 origin,
                                                 id,
                                                 basename,
                                                 flatpak_decomposed_get_ref (ref),
                                                 gpg_data,
                                                 collection_id,
                                                 &created_remote,
                                                 cancellable,
                                                 error);
      if (remote == NULL)
        return NULL;
    }

  if (out_created_remote)
    *out_created_remote = created_remote;

  if (out_ref)
    *out_ref = g_steal_pointer (&ref);

  if (out_checksum)
    *out_checksum = g_steal_pointer (&to_checksum);

  if (out_metadata)
    *out_metadata = g_steal_pointer (&fp_metadata);

  return g_steal_pointer (&remote);
}

Flatpak is a Linux application sandboxing and distribution framework. Prior to versions 1.12.3 and 1.10.6, Flatpak doesn't properly validate that the permissions displayed to the user for an app at install time match the actual permissions granted to the app at runtime, in the case that there's a null byte in the metadata file of an app. Therefore apps can grant themselves permissions without the consent of the user. Flatpak shows permissions to the user during install by reading them from the "xa.metadata" key in the commit metadata. This cannot contain a null terminator, because it is an untrusted GVariant. Flatpak compares these permissions to the *actual* metadata, from the "metadata" file to ensure it wasn't lied to. However, the actual metadata contents are loaded in several places where they are read as simple C-style strings. That means that, if the metadata file includes a null terminator, only the content of the file from *before* the terminator gets compared to xa.metadata. Thus, any permissions that appear in the metadata file after a null terminator are applied at runtime but not shown to the user. So maliciously crafted apps can give themselves hidden permissions. Users who have Flatpaks installed from untrusted sources are at risk in case the Flatpak has a maliciously crafted metadata file, either initially or in an update. This issue is patched in versions 1.12.3 and 1.10.6. As a workaround, users can manually check the permissions of installed apps by checking the metadata file or the xa.metadata key on the commit metadata.


libpq: reject extraneous data after SSL or GSS encryption handshake.

libpq collects up to a bufferload of data whenever it reads data from
the socket.  When SSL or GSS encryption is requested during startup,
any additional data received with the server's yes-or-no reply
remained in the buffer, and would be treated as already-decrypted data
once the encryption handshake completed.  Thus, a man-in-the-middle
with the ability to inject data into the TCP connection could stuff
some cleartext data into the start of a supposedly encryption-protected
database session.

This could probably be abused to inject faked responses to the
client's first few queries, although other details of libpq's behavior
make that harder than it sounds.  A different line of attack is to
exfiltrate the client's password, or other sensitive data that might
be sent early in the session.  That has been shown to be possible with
a server vulnerable to CVE-2021-23214.

To fix, throw a protocol-violation error if the internal buffer
is not empty after the encryption handshake.

Our thanks to Jacob Champion for reporting this problem.

Security: CVE-2021-23222

PQtty(const PGconn *conn)
{
	if (!conn)
		return NULL;
	return "";
}

A man-in-the-middle attacker can inject false responses to the client's first few queries, despite the use of SSL certificate verification and encryption.


conninfo_getval(PQconninfoOption *connOptions,
				const char *keyword)
{
	PQconninfoOption *option;

	option = conninfo_find(connOptions, keyword);

	return option ? option->val : NULL;
}