Searched refs:grad_index (Results 1 – 3 of 3) sorted by relevance
119 GPU_1D_KERNEL_LOOP(grad_index, updates_size) { in SparseApplyAdagradKernel()120 Tindex indices_row = grad_index / col_size; in SparseApplyAdagradKernel()128 Tindex param_index = param_row * col_size + (grad_index % col_size); in SparseApplyAdagradKernel()133 T grad_i = grad[grad_index]; in SparseApplyAdagradKernel()158 GPU_1D_KERNEL_LOOP(grad_index, updates_size) { in SparseApplyProximalAdagradKernel()159 Tindex indices_row = grad_index / col_size; in SparseApplyProximalAdagradKernel()167 Tindex param_index = param_row * col_size + (grad_index % col_size); in SparseApplyProximalAdagradKernel()172 T grad_i = grad[grad_index]; in SparseApplyProximalAdagradKernel()201 GPU_1D_KERNEL_LOOP(grad_index, updates_size) { in SparseApplyFtrlKernel()202 const Tindex indices_row = grad_index / col_size; in SparseApplyFtrlKernel()[all …]
998 for (int grad_index = 0, end = grad.size(); grad_index < end; ++grad_index) { in ForwardpropFromTape() local999 Gradient* grad_tensor = grad[grad_index]; in ForwardpropFromTape()1007 Gradient* in_grad = in_grads[grad_index]; in ForwardpropFromTape()
243 def grad_index(self): member in _GradLoopState