Searched refs:indices_dims (Results 1 – 8 of 8) sorted by relevance
82 std::vector<int64> indices_dims; in FuzzImpl() local88 indices_dims.push_back(dim); in FuzzImpl()90 Tensor indices_tensor(tensorflow::DT_INT32, TensorShape(indices_dims)); in FuzzImpl()106 updates_dims.push_back(indices_dims[i]); in FuzzImpl()107 num_indices *= indices_dims[i]; in FuzzImpl()109 int64 last = indices_dims[indices_rank - 1]; in FuzzImpl()
42 absl::Span<const int64> indices_dims = in XlaScatter() local49 TF_RET_CHECK(!indices_dims.empty()); in XlaScatter()50 num_index_dims = indices_dims.back(); in XlaScatter()58 indices_dims.remove_suffix(1); in XlaScatter()62 for (int64 dim : indices_dims) { in XlaScatter()150 std::vector<int64> expected_updates_dims(indices_dims.begin(), in XlaScatter()151 indices_dims.end()); in XlaScatter()
38 const int indices_dims = indices_shape.dims(); in Compile() local39 const int output_dims = indices_dims + 1; in Compile()56 const int axis = (axis_ == -1) ? indices_dims : axis_; in Compile()
57 const int indices_dims = indices_shape.dims(); in Compute() local58 const int output_dims = indices_dims + 1; in Compute()75 const int axis = (axis_ == -1) ? indices_dims : axis_; in Compute()
45 const int indices_dims = indices->dims->size; in OneHotContext() local46 axis = (params->axis == -1) ? indices_dims : params->axis; in OneHotContext()47 output_dims = indices_dims + 1; in OneHotContext()
1958 const int indices_dims = indices_array.shape().dimensions_count(); in ProcessOneHotOperator() local1959 const int output_dims = indices_dims + 1; in ProcessOneHotOperator()1960 const int axis = op->axis == -1 ? indices_dims : op->axis; in ProcessOneHotOperator()
3438 std::vector<int> indices_dims; in TestConvertGather() member3471 test->AddTestTensor("indices", ok_params[i].indices_dims, 1, in TestConvertGather()
2982 const int indices_dims = indices_shape.DimensionsCount(); in GatherNd() local2983 const int indices_nd = indices_shape.Dims(indices_dims - 1); in GatherNd()2985 for (int i = 0; i < indices_dims - 1; ++i) { in GatherNd()