Method Compile

tensorflow/compiler/tf2xla/kernels/reduction_ops_common.cc:47–123

}

void XlaReductionOp::Compile(XlaOpKernelContext* ctx) {
  const TensorShape data_shape = ctx->InputShape(0);
  const TensorShape axes_tensor_shape = ctx->InputShape(1);
  VLOG(1) << "ReductionOp: " << ctx->op_kernel().name();

  if (axes_tensor_shape.num_elements() == 0) {
    // The reduction axes is an empty vector, which means there are no
    // axes to reduce so just pass the input directly through to the
    // output.
    ctx->SetOutput(0, ctx->Input(0));
    return;
  }

  OP_REQUIRES(ctx, axes_tensor_shape.dims() <= 1,
              errors::InvalidArgument(
                  "Expected scalar or vector as index argument, got ",
                  axes_tensor_shape.DebugString()));

  // Evaluate the constant, reshaping to a 1-vector if it is a scalar.
  std::vector<int64> axes;
  xla::Literal axes_literal;
  OP_REQUIRES_OK(ctx, ctx->ConstantInputReshapedToIntVector(1, &axes));

  VLOG(1) << "data shape: " << data_shape.DebugString();
  VLOG(1) << "axes      : " << absl::StrJoin(axes, ",");

  absl::InlinedVector<bool, 4> bitmap(data_shape.dims(), false);
  std::vector<int64> xla_axes;
  for (int64 i = 0; i < axes_tensor_shape.num_elements(); ++i) {
    int64 index = axes[i];
    OP_REQUIRES(ctx,
                !(index < -data_shape.dims() || index >= data_shape.dims()),
                errors::InvalidArgument("Invalid reduction dimension (", index,
                                        " for input with ", data_shape.dims(),
                                        " dimension(s)"));
    index = (index + data_shape.dims()) % data_shape.dims();
    bitmap[index] = true;
    xla_axes.push_back(index);
  }

  std::vector<int64> final_shape;
  for (int i = 0; i < data_shape.dims(); ++i) {
    if (!bitmap[i]) {
      // If we are not reducing along dimension i.
      int64 dim = data_shape.dim_size(i);
      final_shape.push_back(dim);
    } else if (keep_dims_) {
      // We are reducing along dimension i, but we want to keep the
      // same number of dimensions, so we set the dimension of i to
      // '1'.
      final_shape.push_back(1);
    }
  }

  string desc = ctx->op_kernel().name();

  xla::XlaBuilder* const b = ctx->builder();
  // Construct the builder for the reduction lambda.