Added at::expm1 and at::log1p operators

test/cpp/test_aten_xla_tensor.cpp

@@ -520,6 +520,16 @@ TEST_F(AtenXlaTensorTest, TestExp) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestExpm1) {
+  at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
+  at::Tensor b = at::expm1(a);
+  ForEachDevice([&](const Device& device) {
+    at::Tensor xla_a = bridge::CreateXlaTensor(a, device);
+    at::Tensor xla_b = at::expm1(xla_a);
+    AllClose(b, xla_b, /*rtol=*/1e-3, /*atol=*/1e-5);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestLog) {
   at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
   at::Tensor b = at::log(a);
@@ -530,6 +540,16 @@ TEST_F(AtenXlaTensorTest, TestLog) {
   });
 }
 
+TEST_F(AtenXlaTensorTest, TestLog1p) {
+  at::Tensor a = at::rand({2, 2}, at::TensorOptions(at::kFloat));
+  at::Tensor b = at::log1p(a);
+  ForEachDevice([&](const Device& device) {
+    at::Tensor xla_a = bridge::CreateXlaTensor(a, device);
+    at::Tensor xla_b = at::log1p(xla_a);
+    AllClose(b, xla_b, /*rtol=*/1e-3, /*atol=*/1e-5);
+  });
+}
+
 TEST_F(AtenXlaTensorTest, TestSqrt) {
   at::Tensor a = at::abs(at::rand({2, 2}, at::TensorOptions(at::kFloat)));
   at::Tensor b = at::sqrt(a);

torch_xla/csrc/aten_xla_type.cpp

@@ -158,10 +158,18 @@ at::Tensor AtenXlaType::exp(const at::Tensor& self) const {
   return bridge::AtenFromXlaTensor(XLATensor::exp(bridge::GetXlaTensor(self)));
 }
 
+at::Tensor AtenXlaType::expm1(const at::Tensor& self) const {
+  return bridge::AtenFromXlaTensor(XLATensor::expm1(bridge::GetXlaTensor(self)));
+}
+
 at::Tensor AtenXlaType::log(const at::Tensor& self) const {
   return bridge::AtenFromXlaTensor(XLATensor::log(bridge::GetXlaTensor(self)));
 }
 
+at::Tensor AtenXlaType::log1p(const at::Tensor& self) const {
+  return bridge::AtenFromXlaTensor(XLATensor::log1p(bridge::GetXlaTensor(self)));
+}
+
 at::Tensor AtenXlaType::sqrt(const at::Tensor& self) const {
   return bridge::AtenFromXlaTensor(XLATensor::sqrt(bridge::GetXlaTensor(self)));
 }

torch_xla/csrc/aten_xla_type.h

@@ -43,8 +43,12 @@ class AtenXlaType : public AtenXlaTypeBase {
 
   at::Tensor exp(const at::Tensor& self) const override;
 
+  at::Tensor expm1(const at::Tensor& self) const override;
+
   at::Tensor log(const at::Tensor& self) const override;
 
+  at::Tensor log1p(const at::Tensor& self) const override;
+
   at::Tensor sqrt(const at::Tensor& self) const override;
 
   at::Tensor pow(const at::Tensor& self, at::Scalar exponent) const override;

torch_xla/csrc/ops/ops.cpp

@@ -45,7 +45,9 @@ PTXLA_UNARY_OP(Sinh, at::aten::sinh, xla::Sinh);
 PTXLA_UNARY_OP(Neg, at::aten::neg, xla::Neg);
 PTXLA_UNARY_OP(Abs, at::aten::abs, xla::Abs);
 PTXLA_UNARY_OP(Exp, at::aten::exp, xla::Exp);
+PTXLA_UNARY_OP(Expm1, at::aten::expm1, xla::Expm1);
 PTXLA_UNARY_OP(Log, at::aten::log, xla::Log);
+PTXLA_UNARY_OP(Log1p, at::aten::log1p, xla::Log1p);
 PTXLA_UNARY_OP(Sqrt, at::aten::sqrt, xla::Sqrt);
 
 PTXLA_BINARY_OP(Min, at::aten::min, xla::Min);

torch_xla/csrc/ops/ops.h

@@ -71,8 +71,12 @@ NodePtr Max(const Value& input, const Value& other);
 
 NodePtr Exp(const Value& input);
 
+NodePtr Expm1(const Value& input);
+
 NodePtr Log(const Value& input);
 
+NodePtr Log1p(const Value& input);
+
 NodePtr Sqrt(const Value& input);
 
 NodePtr Pow(const Value& input, const Value& exponent);

torch_xla/csrc/tensor.cpp

@@ -847,10 +847,18 @@ XLATensor XLATensor::exp(const XLATensor& input) {
   return Create(ir::ops::Exp(input.GetIrValue()), input.GetDevice());
 }
 
+XLATensor XLATensor::expm1(const XLATensor& input) {
+  return Create(ir::ops::Expm1(input.GetIrValue()), input.GetDevice());
+}
+
 XLATensor XLATensor::log(const XLATensor& input) {
   return Create(ir::ops::Log(input.GetIrValue()), input.GetDevice());
 }
 
+XLATensor XLATensor::log1p(const XLATensor& input) {
+  return Create(ir::ops::Log1p(input.GetIrValue()), input.GetDevice());
+}
+
 XLATensor XLATensor::sqrt(const XLATensor& input) {
   return Create(ir::ops::Sqrt(input.GetIrValue()), input.GetDevice());
 }

torch_xla/csrc/tensor.h

@@ -251,8 +251,12 @@ class XLATensor {
 
   static XLATensor exp(const XLATensor& input);
 
+  static XLATensor expm1(const XLATensor& input);
+
   static XLATensor log(const XLATensor& input);
 
+  static XLATensor log1p(const XLATensor& input);
+
   static XLATensor sqrt(const XLATensor& input);
 
   static XLATensor pow(const XLATensor& input, at::Scalar exponent);