Method Compile

tensorflow/compiler/tf2xla/kernels/elu_op.cc:62–71

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 one = XlaHelpers::One(b, input_type(0));
    const auto grad = ctx->Input(0);
    const auto activation = ctx->Input(1);
    const auto exp_grad = xla::Mul(grad, xla::Add(activation, one));
    const auto pred = xla::Gt(activation, zero);
    ctx->SetOutput(0, xla::Select(pred, grad, exp_grad));
  }
};

REGISTER_XLA_OP(Name("Elu"), EluOp);