def test_hardshrink(self, dtype):
x = tf.constant([-2.0, -0.5, 0.0, 0.5, 2.0], dtype=dtype)
expected_result = tf.constant([-2.0, 0.0, 0.0, 0.0, 2.0], dtype=dtype)
self.assertAllCloseAccordingToType(_hardshrink_custom_op(x), expected_result)
expected_result = tf.constant([-2.0, 0.0, 0.0, 0.0, 2.0], dtype=dtype)
self.assertAllCloseAccordingToType(
_hardshrink_custom_op(x, lower=-1.0, upper=1.0), expected_result
)
def test_hardshrink(dtype):
x = tf.constant([-2.0, -0.5, 0.0, 0.5, 2.0], dtype=dtype)
expected_result = tf.constant([-2.0, 0.0, 0.0, 0.0, 2.0], dtype=dtype)
test_utils.assert_allclose_according_to_type(_hardshrink_custom_op(x),
expected_result)
expected_result = tf.constant([-2.0, 0.0, 0.0, 0.0, 2.0], dtype=dtype)
test_utils.assert_allclose_according_to_type(
_hardshrink_custom_op(x, lower=-1.0, upper=1.0), expected_result)
def test_invalid():
with pytest.raises(tf.errors.OpError,
match="lower must be less than or equal to upper."):
y = _hardshrink_custom_op(tf.ones(shape=(1, 2, 3)),
lower=2.0,
upper=-2.0)
y.numpy()
def verify_funcs_are_equivalent(self, dtype):
x_np = np.random.uniform(-10, 10, size=(4, 4)).astype(dtype)
x = tf.convert_to_tensor(x_np)
lower = np.random.uniform(-10, 10)
upper = lower + np.random.uniform(0, 10)
with tf.GradientTape(persistent=True) as t:
t.watch(x)
y_native = _hardshrink_custom_op(x, lower, upper)
y_py = _hardshrink_py(x, lower, upper)
self.assertAllCloseAccordingToType(y_native, y_py)
grad_native = t.gradient(y_native, x)
grad_py = t.gradient(y_py, x)
self.assertAllCloseAccordingToType(grad_native, grad_py)
def test_invalid(self):
with self.assertRaisesOpError("lower must be less than or equal to upper."):
y = _hardshrink_custom_op(tf.ones(shape=(1, 2, 3)), lower=2.0, upper=-2.0)
self.evaluate(y)