Add Log1p and DivNoNan gradients
diff --git a/tensorflow/c/eager/gradients_test.cc b/tensorflow/c/eager/gradients_test.cc
index a81d7aa..5dc4f86 100644
--- a/tensorflow/c/eager/gradients_test.cc
+++ b/tensorflow/c/eager/gradients_test.cc
@@ -65,6 +65,8 @@
TF_RETURN_IF_ERROR(registry->Register("Neg", NegRegisterer));
TF_RETURN_IF_ERROR(registry->Register("Sub", SubRegisterer));
TF_RETURN_IF_ERROR(registry->Register("Mul", MulRegisterer));
+ TF_RETURN_IF_ERROR(registry->Register("Log1p", Log1pRegisterer));
+ TF_RETURN_IF_ERROR(registry->Register("DivNoNan", DivNoNanRegisterer));
return Status::OK();
}
@@ -297,6 +299,73 @@
return Status::OK();
}
+// Computes
+// y = log(1 + inputs[0])
+// return grad(y, {inputs[0]})
+Status Log1pGradModel(AbstractContext* ctx,
+ absl::Span<AbstractTensorHandle* const> inputs,
+ absl::Span<AbstractTensorHandle*> outputs,
+ const GradientRegistry& registry) {
+ TapeVSpace vspace(ctx);
+ auto tape = new Tape(/*persistent=*/false);
+ tape->Watch(ToId(inputs[0])); // Watch x.
+ std::vector<AbstractTensorHandle*> log1p_outputs(1);
+ AbstractContextPtr tape_ctx(new TapeContext(ctx, tape, registry));
+ TF_RETURN_IF_ERROR(ops::Log1p(tape_ctx.get(), inputs,
+ absl::MakeSpan(log1p_outputs),
+ "Log1p")); // Compute log(1 + x).
+ std::unordered_map<tensorflow::int64, TapeTensor>
+ source_tensors_that_are_targets;
+
+ std::vector<AbstractTensorHandle*> out_grads;
+ TF_RETURN_IF_ERROR(tape->ComputeGradient(
+ vspace, /*target_tensor_ids=*/{ToId(log1p_outputs[0])},
+ /*source_tensor_ids=*/{ToId(inputs[0])}, source_tensors_that_are_targets,
+ /*output_gradients=*/{}, &out_grads,
+ /*build_default_zeros_grads=*/false));
+ for (auto log1p_output : log1p_outputs) {
+ log1p_output->Unref();
+ }
+ outputs[0] = out_grads[0];
+ delete tape;
+ return Status::OK();
+}
+
+// Computes
+// y = inputs[0] / inputs[1]
+// return grad(y, {inputs[0], inputs[1]})
+Status DivNoNanGradModel(AbstractContext* ctx,
+ absl::Span<AbstractTensorHandle* const> inputs,
+ absl::Span<AbstractTensorHandle*> outputs,
+ const GradientRegistry& registry) {
+ TapeVSpace vspace(ctx);
+ auto tape = new Tape(/*persistent=*/false);
+ tape->Watch(ToId(inputs[0])); // Watch x.
+ tape->Watch(ToId(inputs[1])); // Watch y.
+ std::vector<AbstractTensorHandle*> div_outputs(1);
+ AbstractContextPtr tape_ctx(new TapeContext(ctx, tape, registry));
+ TF_RETURN_IF_ERROR(ops::DivNoNan(tape_ctx.get(), inputs,
+ absl::MakeSpan(div_outputs),
+ "DivNoNan")); // Compute x / y.
+ std::unordered_map<tensorflow::int64, TapeTensor>
+ source_tensors_that_are_targets;
+
+ std::vector<AbstractTensorHandle*> out_grads;
+ TF_RETURN_IF_ERROR(tape->ComputeGradient(
+ vspace, /*target_tensor_ids=*/{ToId(div_outputs[0])},
+ /*source_tensor_ids=*/{ToId(inputs[0]), ToId(inputs[1])},
+ source_tensors_that_are_targets,
+ /*output_gradients=*/{}, &out_grads,
+ /*build_default_zeros_grads=*/false));
+ for (auto div_output : div_outputs) {
+ div_output->Unref();
+ }
+ outputs[0] = out_grads[0];
+ outputs[1] = out_grads[1];
+ delete tape;
+ return Status::OK();
+}
+
AbstractContext* BuildFunction(const char* fn_name) {
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
TF_NewStatus(), TF_DeleteStatus);
@@ -800,6 +869,111 @@
TF_DeleteTensor(result_tensor);
}
+TEST_P(CppGradients, TestLog1pGrad) {
+ std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+ TF_NewStatus(), TF_DeleteStatus);
+ AbstractContextPtr ctx;
+ {
+ AbstractContext* ctx_raw = nullptr;
+ Status s =
+ BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ ctx.reset(ctx_raw);
+ }
+
+ AbstractTensorHandlePtr x;
+ {
+ AbstractTensorHandle* x_raw = nullptr;
+ Status s = TestScalarTensorHandle(ctx.get(), 1.0f, &x_raw);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ x.reset(x_raw);
+ }
+
+ GradientRegistry registry;
+ Status s = RegisterGradients(®istry);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+
+ // Pseudo-code:
+ //
+ // tape.watch(x)
+ // y = log(1 + x)
+ // outputs = tape.gradient(y, x)
+ std::vector<AbstractTensorHandle*> outputs(1);
+ s = RunModel(Log1pGradModel, ctx.get(), {x.get()}, absl::MakeSpan(outputs),
+ /*use_function=*/!std::get<2>(GetParam()), registry);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+
+ TF_Tensor* result_tensor;
+ s = getValue(outputs[0], &result_tensor);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ auto result_value = static_cast<float*>(TF_TensorData(result_tensor));
+ EXPECT_NEAR(*result_value, 0.5, 0.001);
+ outputs[0]->Unref();
+ TF_DeleteTensor(result_tensor);
+ result_tensor = nullptr;
+}
+
+TEST_P(CppGradients, TestDivNoNanGrad) {
+ std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
+ TF_NewStatus(), TF_DeleteStatus);
+ AbstractContextPtr ctx;
+ {
+ AbstractContext* ctx_raw = nullptr;
+ Status s =
+ BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ ctx.reset(ctx_raw);
+ }
+
+ AbstractTensorHandlePtr x;
+ {
+ AbstractTensorHandle* x_raw = nullptr;
+ Status s = TestScalarTensorHandle(ctx.get(), 1.0f, &x_raw);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ x.reset(x_raw);
+ }
+
+ AbstractTensorHandlePtr y;
+ {
+ AbstractTensorHandle* y_raw = nullptr;
+ Status s = TestScalarTensorHandle(ctx.get(), 2.0f, &y_raw);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ y.reset(y_raw);
+ }
+
+ GradientRegistry registry;
+ Status s = RegisterGradients(®istry);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+
+ // Pseudo-code:
+ //
+ // tape.watch(x)
+ // tape.watch(y)
+ // y = x / y
+ // outputs = tape.gradient(y, [x, y])
+ std::vector<AbstractTensorHandle*> outputs(2);
+ s = RunModel(DivNoNanGradModel, ctx.get(), {x.get(), y.get()},
+ absl::MakeSpan(outputs),
+ /*use_function=*/!std::get<2>(GetParam()), registry);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+
+ TF_Tensor* result_tensor;
+ s = getValue(outputs[0], &result_tensor);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ auto result_value = static_cast<float*>(TF_TensorData(result_tensor));
+ EXPECT_NEAR(*result_value, 0.5, 0.001);
+ outputs[0]->Unref();
+ TF_DeleteTensor(result_tensor);
+ result_tensor = nullptr;
+
+ s = getValue(outputs[1], &result_tensor);
+ ASSERT_EQ(errors::OK, s.code()) << s.error_message();
+ result_value = static_cast<float*>(TF_TensorData(result_tensor));
+ EXPECT_NEAR(*result_value, -0.25, 0.001);
+ outputs[1]->Unref();
+ TF_DeleteTensor(result_tensor);
+}
+
TEST_P(CppGradients, TestSetAttrString) {
std::unique_ptr<TF_Status, decltype(&TF_DeleteStatus)> status(
TF_NewStatus(), TF_DeleteStatus);
diff --git a/tensorflow/c/experimental/gradients/math_grad.cc b/tensorflow/c/experimental/gradients/math_grad.cc
index 3ee5294..7b12ff9 100644
--- a/tensorflow/c/experimental/gradients/math_grad.cc
+++ b/tensorflow/c/experimental/gradients/math_grad.cc
@@ -21,10 +21,13 @@
#include "tensorflow/c/experimental/ops/nn_ops.h"
using std::vector;
+using tensorflow::ops::Add;
using tensorflow::ops::Conj;
+using tensorflow::ops::DivNoNan;
using tensorflow::ops::MatMul;
using tensorflow::ops::Mul;
using tensorflow::ops::Neg;
+using tensorflow::ops::OnesLike;
using tensorflow::ops::SqrtGrad;
namespace tensorflow {
@@ -289,6 +292,107 @@
vector<AbstractTensorHandle*> forward_inputs;
};
+class Log1pGradientFunction : public GradientFunction {
+ public:
+ explicit Log1pGradientFunction(vector<AbstractTensorHandle*> f_inputs)
+ : forward_inputs(f_inputs) {}
+
+ Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
+ vector<AbstractTensorHandle*>* grad_outputs) override {
+ /* Given upstream grad U and a Log1p op: Y = log(1 + X), the gradients are:
+ *
+ * dX = U / (1 + X)
+ *
+ */
+
+ AbstractTensorHandle* upstream_grad = grad_inputs[0];
+ AbstractTensorHandle* X = forward_inputs[0];
+
+ grad_outputs->resize(1);
+ vector<AbstractTensorHandle*> temp_outputs(1);
+
+ // Creates Ones
+ std::string name = "OnesLike_Log1p_Grad_X";
+ TF_RETURN_IF_ERROR(
+ OnesLike(ctx->ctx, {X}, absl::MakeSpan(temp_outputs), name.c_str()));
+
+ AbstractTensorHandle* Ones_X = temp_outputs[0];
+
+ // Calculate 1 + X
+ TF_RETURN_IF_ERROR(
+ Add(ctx->ctx, {Ones_X, X}, absl::MakeSpan(temp_outputs), name.c_str()));
+
+ AbstractTensorHandle* XP1 = temp_outputs[0];
+
+ // Calculate U / (1 + X)
+ TF_RETURN_IF_ERROR(DivNoNan(ctx->ctx, {upstream_grad, XP1},
+ absl::MakeSpan(temp_outputs), name.c_str()));
+
+ (*grad_outputs)[0] = temp_outputs[0];
+
+ return Status::OK();
+ }
+ ~Log1pGradientFunction() override {}
+
+ private:
+ vector<AbstractTensorHandle*> forward_inputs;
+};
+
+class DivNoNanGradientFunction : public GradientFunction {
+ public:
+ explicit DivNoNanGradientFunction(vector<AbstractTensorHandle*> f_inputs,
+ vector<AbstractTensorHandle*> f_outputs)
+ : forward_inputs(f_inputs), forward_outputs(f_outputs) {}
+
+ Status Compute(Context* ctx, const IncomingGradients& grad_inputs,
+ vector<AbstractTensorHandle*>* grad_outputs) override {
+ /* Given upstream grad U and a Div op: Z = X/Y, the gradients are:
+ *
+ * dX = U / Y
+ * dY = -U*X / Y^2 = (X/Y) * -U / Y = -U*Z / Y
+ *
+ */
+
+ AbstractTensorHandle* upstream_grad = grad_inputs[0];
+ AbstractTensorHandle* X = forward_inputs[0];
+ AbstractTensorHandle* Y = forward_inputs[1];
+ AbstractTensorHandle* Z = forward_outputs[0];
+
+ grad_outputs->resize(2);
+ vector<AbstractTensorHandle*> temp_outputs(1);
+
+ // Calculate dX = U / Y
+ std::string name = "Div_Grad_X";
+ TF_RETURN_IF_ERROR(DivNoNan(ctx->ctx, {upstream_grad, Y},
+ absl::MakeSpan(temp_outputs), name.c_str()));
+
+ (*grad_outputs)[0] = temp_outputs[0];
+
+ // Calculate dY = -U*Z / Y
+ name = "Neg_Div_Grad_Y";
+ TF_RETURN_IF_ERROR(Neg(ctx->ctx, {upstream_grad},
+ absl::MakeSpan(temp_outputs), name.c_str())); // -U
+ AbstractTensorHandle* MinusU = temp_outputs[0];
+
+ name = "Mul_Div_Grad_Y";
+ TF_RETURN_IF_ERROR(Mul(ctx->ctx, {MinusU, Z}, absl::MakeSpan(temp_outputs),
+ name.c_str())); // -U*Z
+ AbstractTensorHandle* UZ = temp_outputs[0];
+
+ name = "Div_Grad_Y";
+ TF_RETURN_IF_ERROR(DivNoNan(ctx->ctx, {UZ, Y}, absl::MakeSpan(temp_outputs),
+ name.c_str())); // -U*Z / Y
+
+ (*grad_outputs)[1] = temp_outputs[0];
+ return Status::OK();
+ }
+ ~DivNoNanGradientFunction() override {}
+
+ private:
+ vector<AbstractTensorHandle*> forward_inputs;
+ vector<AbstractTensorHandle*> forward_outputs;
+};
+
} // namespace
BackwardFunction* AddRegisterer(const ForwardOperation& op) {
@@ -354,5 +458,23 @@
return new BackwardFunction(gradient_function, default_gradients);
}
+BackwardFunction* Log1pRegisterer(const ForwardOperation& op) {
+ // For ops with a single output, the gradient function is not called if there
+ // is no incoming gradient. So we do not need to worry about creating zeros
+ // grads in this case.
+ auto gradient_function = new Log1pGradientFunction(op.inputs);
+ auto default_gradients = new PassThroughDefaultGradients(op);
+ return new BackwardFunction(gradient_function, default_gradients);
+}
+
+BackwardFunction* DivNoNanRegisterer(const ForwardOperation& op) {
+ // For ops with a single output, the gradient function is not called if there
+ // is no incoming gradient. So we do not need to worry about creating zeros
+ // grads in this case.
+ auto gradient_function = new DivNoNanGradientFunction(op.inputs, op.outputs);
+ auto default_gradients = new PassThroughDefaultGradients(op);
+ return new BackwardFunction(gradient_function, default_gradients);
+}
+
} // namespace gradients
} // namespace tensorflow
diff --git a/tensorflow/c/experimental/gradients/math_grad.h b/tensorflow/c/experimental/gradients/math_grad.h
index d2a0bf2..be73fb5 100644
--- a/tensorflow/c/experimental/gradients/math_grad.h
+++ b/tensorflow/c/experimental/gradients/math_grad.h
@@ -27,6 +27,8 @@
BackwardFunction* NegRegisterer(const ForwardOperation& op);
BackwardFunction* SubRegisterer(const ForwardOperation& op);
BackwardFunction* MulRegisterer(const ForwardOperation& op);
+BackwardFunction* Log1pRegisterer(const ForwardOperation& op);
+BackwardFunction* DivNoNanRegisterer(const ForwardOperation& op);
} // namespace gradients
} // namespace tensorflow
diff --git a/tensorflow/c/experimental/ops/array_ops.cc b/tensorflow/c/experimental/ops/array_ops.cc
index debeba1..1e1820d 100644
--- a/tensorflow/c/experimental/ops/array_ops.cc
+++ b/tensorflow/c/experimental/ops/array_ops.cc
@@ -81,5 +81,17 @@
return op->Execute(outputs, &num_retvals);
}
+Status OnesLike(AbstractContext* ctx,
+ absl::Span<AbstractTensorHandle* const> inputs,
+ absl::Span<AbstractTensorHandle*> outputs, const char* name) {
+ AbstractOperationPtr op(ctx->CreateOperation());
+ TF_RETURN_IF_ERROR(op->Reset("OnesLike", /*raw_device_name=*/nullptr));
+ TF_RETURN_IF_ERROR(MaybeSetOpName(op.get(), name));
+ TF_RETURN_IF_ERROR(op->AddInput(inputs[0]));
+
+ int num_retvals = 1;
+ return op->Execute(outputs, &num_retvals);
+}
+
} // namespace ops
} // namespace tensorflow
diff --git a/tensorflow/c/experimental/ops/array_ops.h b/tensorflow/c/experimental/ops/array_ops.h
index f63412e..035e4de 100644
--- a/tensorflow/c/experimental/ops/array_ops.h
+++ b/tensorflow/c/experimental/ops/array_ops.h
@@ -42,6 +42,10 @@
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs, const char* name);
+Status OnesLike(AbstractContext* ctx,
+ absl::Span<AbstractTensorHandle* const> inputs,
+ absl::Span<AbstractTensorHandle*> outputs, const char* name);
+
} // namespace ops
} // namespace tensorflow
diff --git a/tensorflow/c/experimental/ops/math_ops.cc b/tensorflow/c/experimental/ops/math_ops.cc
index 20aab8a..48b6e63 100644
--- a/tensorflow/c/experimental/ops/math_ops.cc
+++ b/tensorflow/c/experimental/ops/math_ops.cc
@@ -172,5 +172,18 @@
return s;
}
+Status Log1p(AbstractContext* ctx,
+ absl::Span<AbstractTensorHandle* const> inputs,
+ absl::Span<AbstractTensorHandle*> outputs, const char* name) {
+ AbstractOperationPtr log1p_op(ctx->CreateOperation());
+ TF_RETURN_IF_ERROR(log1p_op->Reset("Log1p", /*raw_device_name=*/nullptr));
+ TF_RETURN_IF_ERROR(MaybeSetOpName(log1p_op.get(), name));
+ TF_RETURN_IF_ERROR(log1p_op->AddInput(inputs[0]));
+
+ int num_retvals = 1;
+ Status s = log1p_op->Execute(outputs, &num_retvals);
+ return s;
+}
+
} // namespace ops
} // namespace tensorflow
diff --git a/tensorflow/c/experimental/ops/math_ops.h b/tensorflow/c/experimental/ops/math_ops.h
index 7051e38..4d7d0ae 100644
--- a/tensorflow/c/experimental/ops/math_ops.h
+++ b/tensorflow/c/experimental/ops/math_ops.h
@@ -59,6 +59,10 @@
absl::Span<AbstractTensorHandle* const> inputs,
absl::Span<AbstractTensorHandle*> outputs, const char* name);
+Status Log1p(AbstractContext* ctx,
+ absl::Span<AbstractTensorHandle* const> inputs,
+ absl::Span<AbstractTensorHandle*> outputs, const char* name);
+
} // namespace ops
} // namespace tensorflow
diff --git a/tensorflow/python/framework/experimental/math_ops.cc b/tensorflow/python/framework/experimental/math_ops.cc
index dcb9662..5e988b2 100644
--- a/tensorflow/python/framework/experimental/math_ops.cc
+++ b/tensorflow/python/framework/experimental/math_ops.cc
@@ -86,5 +86,27 @@
ops::Mul(ctx, {a, b}, absl::MakeSpan(outputs), name));
return outputs[0];
});
+ m.def("log1p",
+ [](AbstractContext* ctx, AbstractTensorHandle* a, const char* name) {
+ int num_outputs = 1;
+ std::vector<AbstractTensorHandle*> outputs(1);
+ if (!name) {
+ name = "Log1p";
+ }
+ MaybeRaiseRegisteredFromStatus(
+ ops::Log1p(ctx, {a}, absl::MakeSpan(outputs), name));
+ return outputs[0];
+ });
+ m.def("div_no_nan", [](AbstractContext* ctx, AbstractTensorHandle* a,
+ AbstractTensorHandle* b, const char* name) {
+ int num_outputs = 1;
+ std::vector<AbstractTensorHandle*> outputs(1);
+ if (!name) {
+ name = "DivNoNan";
+ }
+ MaybeRaiseRegisteredFromStatus(
+ ops::DivNoNan(ctx, {a, b}, absl::MakeSpan(outputs), name));
+ return outputs[0];
+ });
}
} // namespace tensorflow
diff --git a/tensorflow/python/framework/experimental/math_ops.py b/tensorflow/python/framework/experimental/math_ops.py
index eee204a..4f3c19f 100644
--- a/tensorflow/python/framework/experimental/math_ops.py
+++ b/tensorflow/python/framework/experimental/math_ops.py
@@ -45,3 +45,13 @@
def mul(a, b, name=None):
ctx = context.get_default()
return _math_ops.mul(ctx, a, b, name)
+
+
+def log1p(a, name=None):
+ ctx = context.get_default()
+ return _math_ops.log1p(ctx, a, name)
+
+
+def div_no_nan(a, b, name=None):
+ ctx = context.get_default()
+ return _math_ops.div_no_nan(ctx, a, b, name)
diff --git a/tensorflow/python/framework/experimental/tape.cc b/tensorflow/python/framework/experimental/tape.cc
index 8e7bcfa..a29d7cc 100644
--- a/tensorflow/python/framework/experimental/tape.cc
+++ b/tensorflow/python/framework/experimental/tape.cc
@@ -39,6 +39,8 @@
TF_RETURN_IF_ERROR(registry->Register("Neg", NegRegisterer));
TF_RETURN_IF_ERROR(registry->Register("Sub", SubRegisterer));
TF_RETURN_IF_ERROR(registry->Register("Mul", MulRegisterer));
+ TF_RETURN_IF_ERROR(registry->Register("Log1p", Log1pRegisterer));
+ TF_RETURN_IF_ERROR(registry->Register("DivNoNan", DivNoNanRegisterer));
return Status::OK();
}
diff --git a/tensorflow/python/framework/experimental/unified_api_test.py b/tensorflow/python/framework/experimental/unified_api_test.py
index 98615d4..b4b9dbe 100644
--- a/tensorflow/python/framework/experimental/unified_api_test.py
+++ b/tensorflow/python/framework/experimental/unified_api_test.py
@@ -305,6 +305,99 @@
self.assertAllEqual(eager_outputs[0].numpy(), [3., 4.])
self.assertAllEqual(eager_outputs[1].numpy(), [1., 2.])
+ @parameterized.named_parameters([
+ ("Graph", False),
+ ("Mlir", True),
+ ])
+ def testLog1p(self, use_mlir):
+ if use_mlir:
+ SetTracingImplementation("mlir")
+
+ def model(a):
+ return unified_math_ops.log1p(a)
+
+ with context_lib.set_default(get_immediate_execution_context()):
+ a = TensorCastHelper(constant_op.constant([1.]))
+
+ func_output = def_function.function(model)(a)
+ self.assertArrayNear(func_output.numpy(), [0.69314], 0.001)
+
+ eager_output = model(a)
+ self.assertArrayNear(eager_output.numpy(), [0.69314], 0.001)
+
+ @parameterized.named_parameters([
+ ("Graph", False),
+ ("Mlir", True),
+ ])
+ def testLog1pGrad(self, use_mlir):
+ if use_mlir:
+ SetTracingImplementation("mlir")
+
+ def model(a):
+ with tape_lib.GradientTape() as tape:
+ tape.watch(a)
+ result = unified_math_ops.log1p(a)
+ grads = tape.gradient(result, a)
+ return grads
+
+ with context_lib.set_default(get_immediate_execution_context()):
+ a = TensorCastHelper(constant_op.constant([1.]))
+
+ func_outputs = def_function.function(model)(a)
+ self.assertArrayNear(func_outputs.numpy(), [0.5], 0.001)
+
+ eager_outputs = model(a)
+ self.assertArrayNear(eager_outputs.numpy(), [0.5], 0.001)
+
+ @parameterized.named_parameters([
+ ("Graph", False),
+ ("Mlir", True),
+ ])
+ def testDivNoNan(self, use_mlir):
+ if use_mlir:
+ SetTracingImplementation("mlir")
+
+ def model(a, b):
+ return unified_math_ops.div_no_nan(a, b)
+
+ with context_lib.set_default(get_immediate_execution_context()):
+ a = TensorCastHelper(constant_op.constant([2.]))
+ b = TensorCastHelper(constant_op.constant([4.]))
+
+ func_output = def_function.function(model)(a, b)
+ self.assertArrayNear(func_output.numpy(), [0.5], 0.001)
+
+ eager_output = model(a, b)
+ self.assertArrayNear(eager_output.numpy(), [0.5], 0.001)
+
+ @parameterized.named_parameters([
+ ("Graph", False),
+ ("Mlir", True),
+ ])
+ def testDivNoNanGrad(self, use_mlir):
+ if use_mlir:
+ SetTracingImplementation("mlir")
+
+ def model(a, b):
+ with tape_lib.GradientTape() as tape:
+ tape.watch(a)
+ tape.watch(b)
+ result = unified_math_ops.div_no_nan(a, b)
+ grads = tape.gradient(result, [a, b])
+ return grads
+
+ with context_lib.set_default(get_immediate_execution_context()):
+ a = TensorCastHelper(constant_op.constant([2.]))
+ b = TensorCastHelper(constant_op.constant([4.]))
+
+ func_outputs = def_function.function(model)(a, b)
+ self.assertArrayNear(func_outputs[0].numpy(), [0.25], 0.001)
+ self.assertArrayNear(func_outputs[1].numpy(), [-0.125], 0.001)
+
+ eager_outputs = model(a, b)
+ self.assertArrayNear(eager_outputs[0].numpy(), [0.25], 0.001)
+ self.assertArrayNear(eager_outputs[1].numpy(), [-0.125], 0.001)
+
class UnifiedTapeBenchmark(test.Benchmark):
