Method Compile

tensorflow/compiler/tf2xla/kernels/softmax_op.cc:42–89

  }

  void Compile(XlaOpKernelContext* ctx) override {
    const TensorShape logits_shape = ctx->InputShape(0);
    OP_REQUIRES(ctx, TensorShapeUtils::IsVectorOrHigher(logits_shape),
                errors::InvalidArgument("logits must have >= 1 dimension, got ",
                                        logits_shape.DebugString()));

    // Major dimensions are batch dimensions, minor dimension is the class
    // dimension.
    std::vector<int64> batch_dims(logits_shape.dims() - 1);
    std::iota(batch_dims.begin(), batch_dims.end(), 0);
    const int kClassDim = logits_shape.dims() - 1;

    const DataType type = input_type(0);
    const xla::PrimitiveType xla_type = ctx->input_xla_type(0);
    auto logits = ctx->Input(0);
    xla::XlaOp softmax;
    if (is_on_gpu_) {
      softmax = xla::Softmax(logits, kClassDim, log_);
    } else {
      xla::XlaBuilder* const b = ctx->builder();
      const xla::XlaComputation& max_func = *ctx->GetOrCreateMax(type);

      // Find the max in each batch, resulting in a tensor of shape [batch]
      auto logits_max =
          xla::Reduce(logits, xla::MinValue(b, xla_type), max_func, {kClassDim});
      // Subtract the max in batch b from every element in batch b. Broadcasts
      // along the batch dimension.
      auto shifted_logits = xla::Sub(logits, logits_max, batch_dims);
      auto exp_shifted = xla::Exp(shifted_logits);
      const DataType accumulation_type = XlaHelpers::SumAccumulationType(type);
      xla::PrimitiveType xla_accumulation_type;
      OP_REQUIRES_OK(ctx, DataTypeToPrimitiveType(accumulation_type,
                                                  &xla_accumulation_type));
      auto converted =
          xla::ConvertElementType(exp_shifted, xla_accumulation_type);
      auto reduce =
          xla::Reduce(converted, xla::Zero(b, xla_accumulation_type),
                      *ctx->GetOrCreateAdd(accumulation_type), {kClassDim});
      auto sum = XlaHelpers::ConvertElementType(reduce, type);
      softmax =
          log_
              // softmax = shifted_logits - log(sum(exp(shifted_logits)))
              ? xla::Sub(shifted_logits, xla::Log(sum), batch_dims)
              // softmax = exp(shifted_logits) / sum(exp(shifted_logits))
              : xla::Div(exp_shifted, sum, batch_dims);
    }
    ctx->SetOutput(0, softmax);
  }

 private:
  bool log_;