tfp.substrates.jax.math.clip_by_value_preserve_gradient

Clips values to a specified min and max while leaving gradient unaltered.

tfp.substrates.jax.math.clip_by_value_preserve_gradient(
    t, clip_value_min, clip_value_max, name=None
)

Like tf.clip_by_value, this function returns a tensor of the same type and shape as input t but with values clamped to be no smaller than to clip_value_min and no larger than clip_value_max. Unlike tf.clip_by_value, the gradient is unaffected by this op, i.e.,

tf.gradients(tfp.math.clip_by_value_preserve_gradient(x), x)[0]
# ==> ones_like(x)

t A Tensor.
clip_value_min A scalar Tensor, or a Tensor with the same shape as t. The minimum value to clip by.
clip_value_max A scalar Tensor, or a Tensor with the same shape as t. The maximum value to clip by.
name A name for the operation (optional). Default value: 'clip_by_value_preserve_gradient'.

clipped_t A clipped Tensor.