def renorm(x, axis, max_norm):
    """Renormalizes the sub-tensors along axis such that they do not exceed norm max_norm."""
rank = tf.rank(x)
bigrange = tf.range(-1, rank + 1)
dims = tf.slice(tf.concat(0, [tf.slice(bigrange, [0], [1 + axis]), tf.slice
    (bigrange, [axis + 2], [-1])]), [1], rank - [1])
l2norm_inv = tf.rsqrt(tf.reduce_sum(x * x, dims, keep_dims=True))
scale = max_norm * tf.minimum(l2norm_inv, tf.constant(1.0 / max_norm))
return tf.mul(scale, x)
Renormalize weight matrix using TensorFlow