Searched refs:mat_g_size (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/contrib/opt/python/training/ |
| D | shampoo.py | 231 def _compute_power_svd(self, var, mat_g, mat_g_size, alpha, mat_h_slot_name): argument
247 if mat_g_size == 1:
251 math_ops.cast(mat_g_size, dtypes.int32))
260 def _compute_power_iter(self, var, mat_g, mat_g_size, alpha, mat_h_slot_name, argument
264 mat_g_sqrt = matrix_functions.matrix_square_root(mat_g, mat_g_size,
268 mat_g_size,
278 def _compute_power(self, var, mat_g, mat_g_size, alpha, mat_h_slot_name=None): argument
282 return self._compute_power_iter(var, mat_g, mat_g_size, alpha,
285 return self._compute_power_svd(var, mat_g, mat_g_size, alpha,
|
| D | matrix_functions.py | 77 mat_g_size, argument
104 identity = linalg_ops.eye(math_ops.cast(mat_g_size, dtypes.int32))
136 if mat_g_size == 1:
|