Method Compile

tensorflow/compiler/tf2xla/kernels/elu_op.cc:91–104

Return the lhs (incoming gradient) if the rhs (input feature) > 0, otherwise return lhs * (1 + rhs).

  // Return the lhs (incoming gradient) if the rhs (input feature) > 0,
  // otherwise return lhs * (1 + rhs).
  void Compile(XlaOpKernelContext* ctx) override {
    xla::XlaBuilder* b = ctx->builder();
    const auto zero = XlaHelpers::Zero(b, input_type(0));
    const auto scale = XlaHelpers::FloatLiteral(b, input_type(0),
            1.0507009873554804934193349852946);
    const auto scale_alpha = XlaHelpers::FloatLiteral(b, input_type(0),
            1.7580993408473768599402175208123);
    const auto grad = ctx->Input(0);
    const auto activation = ctx->Input(1);
    const auto lin_grad = xla::Mul(grad, scale);
    const auto exp_grad = xla::Mul(grad, xla::Add(activation, scale_alpha));
    const auto pred = xla::Gt(activation, zero);
    ctx->SetOutput(0, xla::Select(pred, lin_grad, exp_grad));
  }
};

REGISTER_XLA_OP(Name("Selu"), SeluOp);