Initial version of cumprod #6211

iislucas · 2022-03-10T11:29:07Z

Initial implementation of cumprod:

Addresses: #6210

Includes implementation for GPU/WebGL, WASM, and CPU.

Does not (yet) include gradient computation (that's quite a bit harder).

I did this based on replicating files for cumsum. If there are more implementations that use the same scanning structure, it might make sense to abstract away some of the common code. But currently there's only two functions, so probably best to wait for a third to avoid premature generalisation.

Tested locally with:

$ yarn run bazel test --test_output=streamed :tfjs-core_test :test_snippets_test :tfjs-core_node_test :tfjs-core_async_backends_test src:worker_node_test src:worker_test --//:grep=cumprod

This change is

pyu10055

Thank you for contributing the cumprod op, the gradient function can come later, for reference the similar grad def for cumsum is in src/gradients/Cumsum_grad.ts

I agree your suggestion on refactoring the scan ops later, that the struct of scan ops are very similar, the similar approach we had is unary op.

The PR looks good overall, just couple minor nitpick. Thanks!

Reviewable status: 0 of 1 approvals obtained (waiting on @iislucas)

tfjs-backend-cpu/src/kernels/Cumprod.ts, line 3 at r1 (raw file):

/**
 * @license
 * Copyright 2020 Google LLC. All Rights Reserved.

2022

tfjs-backend-webgl/src/cumprod_gpu.ts, line 33 at r1 (raw file):

    let condition = '';
    let idxString = '';
    // When exclusive is set, the cumsum op becomes roll op that copies the

cumprod

tfjs-backend-webgl/src/cumprod_gpu.ts, line 71 at r1 (raw file):

    return `${name}.x, ${name}.y, ${name}.z, ${name}.w`;
  } else {
    throw Error(`Cumulative sum for rank ${rank} is not yet supported`);

prod

tfjs-backend-webgl/src/cumprod_gpu.ts, line 85 at r1 (raw file):

    return `${name}.w`;
  } else {
    throw Error(`Cumulative sum for rank ${rank} is not yet supported`);

prod

tfjs-backend-webgl/src/kernels/Cumprod.ts, line 63 at r1 (raw file):

    backend.disposeIntermediateTensorInfo(prevResult);
  }
  // For exclusive cumsum, shift the end result in the direction of sum

update this comment according to the cumprod definition.

tfjs-core/src/ops/cumprod.ts, line 40 at r1 (raw file):

 * ```
 *
 * @param x The input tensor to be summed.

cumulative multiplied.

tfjs-core/src/ops/cumprod_test.ts, line 41 at r1 (raw file):

    const res = tf.tensor1d([1, 2, 3, 4]).cumprod(0, exclusive);
    expect(res.shape).toEqual([4]);
    expectArraysClose(await res.data(), [0, 1, 2, 6]);

1,1,2,6?

tfjs-core/src/ops/cumprod_test.ts, line 49 at r1 (raw file):

    const res = tf.tensor1d([1, 2, 3, 4]).cumprod(0, exclusive, reverse);
    expect(res.shape).toEqual([4]);
    expectArraysClose(await res.data(), [6, 2, 1, 0]);

6,2,1,1?

tfjs-core/src/ops/cumprod_test.ts, line 52 at r1 (raw file):

  });

  it('gradient: 1D', async () => {

you might want to register the gradient function similar to cumsumGradConfig in src/gradients/Cumsum_grad.ts

tfjs-core/src/ops/cumprod_test.ts, line 78 at r1 (raw file):

      .cumprod(1);
    expect(res.shape).toEqual([2, 2]);
    expectArraysClose(await res.data(), [1, 3, 3, 7]);

I think all of these test results need to be updated for cumprod

iislucas

Thanks for looking over, tests fixed and did a careful search to catch all missing "sum" -> "product" language changes.

Reviewable status: 0 of 1 approvals obtained (waiting on @pyu10055)

tfjs-backend-cpu/src/kernels/Cumprod.ts, line 3 at r1 (raw file):