| 31 | namespace tensorflow { |
| 32 | |
| 33 | xla::StatusOr<xla::XlaOp> XlaScatter( |
| 34 | const xla::XlaOp& buffer, const xla::XlaOp& updates, |
| 35 | const xla::XlaOp& indices, bool indices_are_vectors, |
| 36 | const std::function<xla::XlaOp(xla::XlaOp, xla::XlaOp, xla::XlaBuilder*)>& |
| 37 | combiner, |
| 38 | xla::XlaBuilder* builder) { |
| 39 | TF_ASSIGN_OR_RETURN(xla::Shape buffer_shape, builder->GetShape(buffer)); |
| 40 | TF_ASSIGN_OR_RETURN(xla::Shape updates_shape, builder->GetShape(updates)); |
| 41 | TF_ASSIGN_OR_RETURN(xla::Shape indices_shape, builder->GetShape(indices)); |
| 42 | absl::Span<const int64> indices_dims = |
| 43 | xla::AsInt64Slice(indices_shape.dimensions()); |
| 44 | |
| 45 | // If the indices are N-dimensional, the minor dimension of indices contains |
| 46 | // the indices to update. Otherwise the indices are all scalars. |
| 47 | int64 num_index_dims = 1; |
| 48 | if (indices_are_vectors) { |
| 49 | TF_RET_CHECK(!indices_dims.empty()); |
| 50 | num_index_dims = indices_dims.back(); |
| 51 | if (num_index_dims > buffer_shape.rank()) { |
| 52 | return errors::InvalidArgument( |
| 53 | "The size of the minor dimension of the indices (shape: ", |
| 54 | xla::ShapeUtil::HumanString(indices_shape), |
| 55 | ") must be <= the rank of the buffer (shape: ", |
| 56 | xla::ShapeUtil::HumanString(buffer_shape), ")"); |
| 57 | } |
| 58 | indices_dims.remove_suffix(1); |
| 59 | } |
| 60 | |
| 61 | int64 num_indices = 1; |
| 62 | for (int64 dim : indices_dims) { |
| 63 | num_indices *= dim; |
| 64 | } |
| 65 | |
| 66 | // Degenerate case: nothing to update. Return the buffer unchanged. |
| 67 | if (num_indices == 0) { |
| 68 | return buffer; |
| 69 | } |
| 70 | |
| 71 | // If any of the indexed dimensions are zero in the buffer, the update cannot |
| 72 | // succeed since it updates a slice of size 1. |
| 73 | for (int64 i = 0; i < num_index_dims; ++i) { |
| 74 | if (xla::ShapeUtil::GetDimension(buffer_shape, i) == 0) { |
| 75 | return errors::InvalidArgument("Scatter dimension ", i, |
| 76 | " is of size zero in tensor with shape ", |
| 77 | xla::ShapeUtil::HumanString(buffer_shape)); |
| 78 | } |
| 79 | } |
| 80 | |
| 81 | // Example of a 1-D scatter that updates two [3,1] tensors in a tensor of |
| 82 | // shape [3,3]: |
| 83 | // NOTE: ***This case will not be generated by any of the tf.scatter ops.*** |
| 84 | // |
| 85 | // operand = s32[3,3] parameter(0) |
| 86 | // indices = s32[2] parameter(1) |
| 87 | // updates = s32[3,2] parameter(2) |
| 88 | // scatter = s32[3,3] scatter(operand, indices, updates), |
| 89 | // to_apply=update_computation, |
| 90 | // update_window_dims={0}, |
no test coverage detected