Function sparseSegmentSum_

tfjs-core/src/ops/sparse/sparse_segment_sum.ts:59–88

* Computes the sum along sparse segments of a tensor. * * ```js * const c = tf.tensor2d([[1,2,3,4], [-1,-2,-3,-4], [5,6,7,8]]); * // Select two rows, one segment. * const result1 = tf.sparse.sparseSegmentSum(c, * tf.tensor1d([0, 1], 'int32'), *

(
    data: Tensor|TensorLike, indices: Tensor1D|TensorLike,
    segmentIds: Tensor1D|TensorLike)

 * @doc {heading: 'Operations', subheading: 'Sparse'}
 */
function sparseSegmentSum_(
    data: Tensor|TensorLike, indices: Tensor1D|TensorLike,
    segmentIds: Tensor1D|TensorLike): Tensor {
  const $data = convertToTensor(data, 'data', 'sparseSegmentSum');
  const $indices =
      convertToTensor(indices, 'indices', 'sparseSegmentSum', 'int32');
  const $segmentIds =
      convertToTensor(segmentIds, 'segmentIds', 'sparseSegmentSum', 'int32');

  if ($data.rank < 1) {
    throw new Error(
        `Data should be at least 1 dimensional but received scalar`);
  }
  if ($indices.rank !== 1) {
    throw new Error(`Indices should be Tensor1D but received shape
         ${$indices.shape}`);
  }
  if ($segmentIds.rank !== 1) {
    throw new Error(`Segment ids should be Tensor1D but received shape
         ${$segmentIds.shape}`);
  }

  const inputs: SparseSegmentSumInputs = {
    data: $data,
    indices: $indices,
    segmentIds: $segmentIds
  };

  return ENGINE.runKernel(SparseSegmentSum, inputs as {});
}

export const sparseSegmentSum = /* @__PURE__ */ op({sparseSegmentSum_});