Searched refs:indices_dims (Results 1 – 7 of 7) 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()
58 const int indices_dims = indices_shape.dims(); in Compute() local59 const int output_dims = indices_dims + 1; in Compute()76 const int axis = (axis_ == -1) ? indices_dims : axis_; in Compute()
47 const int indices_dims = indices->dims->size; in OneHotContext() local48 axis = (params->axis == -1) ? indices_dims : params->axis; in OneHotContext()49 output_dims = indices_dims + 1; in OneHotContext()
1073 const int indices_dims = indices_shape.DimensionsCount(); in GatherNdHelper() local1074 ret.indices_nd = indices_shape.Dims(indices_dims - 1); in GatherNdHelper()1076 for (int i = 0; i < indices_dims - 1; ++i) { in GatherNdHelper()
1981 const int indices_dims = indices_array.shape().dimensions_count(); in ProcessOneHotOperator() local1982 const int output_dims = indices_dims + 1; in ProcessOneHotOperator()1983 const int axis = op->axis == -1 ? indices_dims : op->axis; in ProcessOneHotOperator()