Function sparseSegmentMean_

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

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

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

 * @doc {heading: 'Operations', subheading: 'Sparse'}
 */
function sparseSegmentMean_(
    data: Tensor|TensorLike, indices: Tensor1D|TensorLike,
    segmentIds: Tensor1D|TensorLike): Tensor {
  const $data = convertToTensor(data, 'data', 'sparseSegmentMean');
  const $indices =
      convertToTensor(indices, 'indices', 'sparseSegmentMean', 'int32');
  const $segmentIds =
      convertToTensor(segmentIds, 'segmentIds', 'sparseSegmentMean', '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: SparseSegmentMeanInputs = {
    data: $data,
    indices: $indices,
    segmentIds: $segmentIds
  };

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

export const sparseSegmentMean = /* @__PURE__ */ op({sparseSegmentMean_});