Method Compile

tensorflow/compiler/tf2xla/kernels/bias_ops.cc:90–119

  }

  void Compile(XlaOpKernelContext* ctx) override {
    const TensorShape out_backprop_shape = ctx->InputShape(0);

    OP_REQUIRES(ctx, TensorShapeUtils::IsMatrixOrHigher(out_backprop_shape),
                errors::InvalidArgument("Input tensor must be at least 2D: ",
                                        out_backprop_shape.DebugString()));

    // feature_dim is the channel (C) dimension of the data.
    int feature_dim = (data_format_ == FORMAT_NHWC)
                          ? out_backprop_shape.dims() - 1
                          : /*data_format == FORMAT_NCHW*/ 1;
    OP_REQUIRES(
        ctx, feature_dim >= 0,
        errors::InvalidArgument("Input tensor does not have enough dimensions "
                                "to contain the feature dimension"));

    std::vector<int64> reduce_dims(out_backprop_shape.dims() - 1);
    std::iota(reduce_dims.begin(), reduce_dims.begin() + feature_dim, 0);
    std::iota(reduce_dims.begin() + feature_dim, reduce_dims.end(),
              feature_dim + 1);
    xla::XlaBuilder* const b = ctx->builder();
    const DataType accumulation_type =
        XlaHelpers::SumAccumulationType(input_type(0));
    auto converted =
        XlaHelpers::ConvertElementType(ctx->Input(0), accumulation_type);
    auto reduce =
        xla::Reduce(converted, XlaHelpers::Zero(b, accumulation_type),
                    *ctx->GetOrCreateAdd(accumulation_type), reduce_dims);
    ctx->SetOutput(0, XlaHelpers::ConvertElementType(reduce, input_type(0)));
  }

 private:
  TensorFormat data_format_;