Function createTensorFromGPUTest

tfjs-backend-webgpu/src/backend_webgpu_test.ts:562–680

(zeroCopy = false)

}

function createTensorFromGPUTest(zeroCopy = false) {
  it('use default shape and data type(float32)', async () => {
    await testCreateTensorFromGPUBuffer('float32', true, zeroCopy);
  });

  it('work for float32', async () => {
    await testCreateTensorFromGPUBuffer('float32', false, zeroCopy);
  });

  it('work for int32', async () => {
    await testCreateTensorFromGPUBuffer('int32', false, zeroCopy);
  });

  it('work for read', async () => {
    const webGPUBackend = tf.backend() as WebGPUBackend;
    const device = webGPUBackend.device;
    const aData = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
    const dtype = 'float32';
    const aBuffer = createGPUBufferFromData(device, aData, dtype);
    const shape: number[] = [aData.length];
    const a = tf.tensor({buffer: aBuffer, zeroCopy}, shape, dtype);
    if (zeroCopy !== true) {
      aBuffer.destroy();
    }
    await a.data();
    if (zeroCopy === true) {
      aBuffer.destroy();
    }
  });

  it('two tensors share the same GPUBuffer', async () => {
    const webGPUBackend = tf.backend() as WebGPUBackend;
    const device = webGPUBackend.device;
    const aData = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
    const dtype = 'float32';
    const aBuffer = createGPUBufferFromData(device, aData, dtype);
    const startNumBytes = tf.memory().numBytes;
    const startNumTensors = tf.memory().numTensors;
    const shape: number[] = [aData.length];
    const webGPUData = {buffer: aBuffer, zeroCopy};
    const a = tf.tensor(webGPUData, shape, dtype);
    const b = tf.tensor(webGPUData, shape, dtype);
    if (zeroCopy !== true) {
      aBuffer.destroy();
    }
    const result = tf.add(a, b);
    const expected =
        [2, 4, 6, 8, 10, 12, 14, 16, 18, 20, 22, 24, 26, 28, 30, 32];
    tf.test_util.expectArraysClose(await result.data(), expected);
    a.dispose();
    b.dispose();
    result.dispose();
    const endNumBytes = tf.memory().numBytes;
    const endNumTensors = tf.memory().numTensors;
    expect(endNumBytes - startNumBytes).toEqual(0);
    expect(endNumTensors - startNumTensors).toEqual(0);
    if (zeroCopy === true) {
      aBuffer.destroy();