custom bundle tool

[tafsiri]

Introduces a CLI tool in the tfjs package, tfjs-custom-bundle that can be used to create entry files (modules) that can be tree shaken by a bundler as part of whole program optimization. The tool isn't complete but this is the first part and in order to keep PRs smaller I figured we could review it piece by piece. Currently it only supports tree shaking programs that only use tfjs-core.

It takes a config file as input: example1, example2
And produces an ES module as output that replaces the import * as tf from ... statements in a program: example1, example2.

Note: eventually the examples in that repo linked above would become e2e tests in this repo. But it is sufficiently experimental at this stage that I want to keep them outside.

You can see the results of tree shaking via rollup: example1, (for example 2 not all those kernels have been modularized in those backends, so while the properly tree shake bundle is created it is incomplete and can't run). In particular look at the "kernels" subfolder of each backend, only the modular kernels specified in the config file are included.

To achieve this we also need to tell bundlers like rollup and webpack which files are sideffect free. This involves some file reorganization and importantly of a list of files in package.json for files that do have global side effects. Please review these and point out if there any missing (anything that sets up a registry or auto registers anything would be included).

FYI running this does require some manual steps as we can't yet remove register_all_chained_ops from core without breaking api. This is planned for 3.x but we will work around this in the mean time.

As a note the user uses this tool by installing the @tensorflow/tfjs package and then running

npx tfjs-custom-bundle --config path_to_config.json the rest is as described in the design doc.

To see the logs from the Cloud Build CI, please join either our discussion or announcement mailing list.

---

This change is 

[annxingyuan]

This is super exciting!! I couldn't think of other side-effectful files...

Reviewable status: 0 of 1 approvals obtained (waiting on @annxingyuan, @lina128, and @tafsiri)

---

tfjs/tools/custom_bundle/cli.ts, line 68 at r1 (raw file):

  console.log(`Using custom bundle configuration from ${
      configFilePath}. Final config:`);
  console.log(`${JSON.stringify(config, null, 2)}\n`);

could we name null and 2

---

tfjs/tools/custom_bundle/custom_module_test.ts, line 46 at r1 (raw file):

    expect(res).toContain('import BACKEND FastBcknd');
    expect(res).toContain('import KERNEL MathKrnl from BACKEND FastBcknd');
    expect(res).toContain('registerKernel(MathKrnl_FastBcknd)');

would it make sense to assert what the entire module string should be? might be helpful to see what the full output would be.

---

tfjs/tools/custom_bundle/custom_module_test.ts, line 138 at r1 (raw file):

    const gradIndex = res.indexOf('GRADIENT');
    expect(res.indexOf('GRADIENT', gradIndex + 1))

could use str.includes for readability

---

tfjs/tools/custom_bundle/esm_module_provider_test.ts, line 31 at r1 (raw file):

  });

  // tslint:disable-next-line: ban

did you mean to include this?

[lina128]

+1 Super exciting! Thank you Yannick!!!

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @lina128 and @tafsiri)

---

tfjs-backend-cpu/src/base.ts, line 20 at r1 (raw file):

/*
 * base.ts contains all the exports from tfjs-backend-cpu
 * but skips auto-kernel registration

sugg: without auto-kernel registration

[tafsiri]

Reviewable status: complete! 1 of 1 approvals obtained (waiting on @annxingyuan and @lina128)

---

tfjs/tools/custom_bundle/cli.ts, line 68 at r1 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

could we name null and 2

Done

---

tfjs/tools/custom_bundle/custom_module_test.ts, line 46 at r1 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

would it make sense to assert what the entire module string should be? might be helpful to see what the full output would be.

I think that would be brittle to white space changes and the like. These assertions don't care about the order these appear in just that they are correctly constructed given the passed in provider. e2e tests down the line will test that the file is valid and builds a bundle that runs the program correctly and will be an overall more robust check.

---

tfjs/tools/custom_bundle/custom_module_test.ts, line 138 at r1 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

could use str.includes for readability

Here I'm specifically checking that the substring doesn't appear twice, indexOf allows me to start after the first appearance of GRADIENT without having to split the string.

---

tfjs/tools/custom_bundle/esm_module_provider_test.ts, line 31 at r1 (raw file):

Previously, annxingyuan (Ann Yuan) wrote…

did you mean to include this?

Nope. Removed.

---

tfjs-backend-cpu/src/base.ts, line 20 at r1 (raw file):

Previously, lina128 (Na Li) wrote…

sugg: without auto-kernel registration

done