Searched refs:grad_idx (Results 1 – 2 of 2) sorted by relevance
161 const Tensor* grad_idx = std::get<0>(*grad); in AllocateAndAssignToAccumGradFunction() local164 const int64 nnz = grad_idx->dim_size(0); in AllocateAndAssignToAccumGradFunction()171 accum_idx_vec_->push_back(grad_idx->vec<int64>()(i)); in AllocateAndAssignToAccumGradFunction()197 const Tensor* grad_idx = std::get<0>(*grad); in AddToAccumGradFunction() local201 const int64 grad_nnz = grad_idx->dim_size(0); in AddToAccumGradFunction()219 switch (cmp(accum_idx_vec_, grad_idx, i, j)) { in AddToAccumGradFunction()297 sum_indices_vec->push_back(grad_idx->vec<int64>()(idx_b)); in AddToAccumGradFunction()
38 int grad_idx, ShapeHandle* s) { in HandleGradAndIndicesInputs() argument39 ShapeHandle grad = ShapeOrHandleShape(c, grad_idx); in HandleGradAndIndicesInputs()46 TF_RETURN_IF_ERROR(c->WithRank(c->input(grad_idx + 1), 1, &indices)); in HandleGradAndIndicesInputs()