(self, a, return_inverse=False)
| 3100 | return tnp.diag(a, k) |
| 3101 | |
| 3102 | def unique(self, a, return_inverse=False): |
| 3103 | y, idx = tf.unique(tf.reshape(a, [-1])) |
| 3104 | sort_idx = tf.argsort(y) |
| 3105 | y_prime = tf.gather(y, sort_idx) |
| 3106 | if return_inverse: |
| 3107 | inv_sort_idx = tf.math.invert_permutation(sort_idx) |
| 3108 | return y_prime, tf.gather(inv_sort_idx, idx) |
| 3109 | else: |
| 3110 | return y_prime |
| 3111 | |
| 3112 | def logsumexp(self, a, axis=None): |
| 3113 | return tf.math.reduce_logsumexp(a, axis=axis) |