[webgl] Modularize sqrt

[mattsoulanille]

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

---

This change is 

[tafsiri]

Reviewed 2 of 3 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @mattsoulanille)

---

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

  // TODO(msoulanille) remove this once the backend has been modularized
  // a copy is needed here because backend_webgl tests use sqrt.

Those tests should still work if this is removed, they call sqrt via the op in core which should dispatch to the modular sqrt kernel you added below. What error are you seeing when you remove this?

[mattsoulanille]

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

---

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

Previously, tafsiri (Yannick Assogba) wrote…

Those tests should still work if this is removed, they call sqrt via the op in core which should dispatch to the modular sqrt kernel you added below. What error are you seeing when you remove this?

I omitted sqrt from a previous PR because the backend_webgl tests were failing. I forget what the error was, so I've pushed a commit removing sqrt from backend_webgl so we can see if it appears again.

[mattsoulanille]

Looks like sqrt is not being loaded before the tests are run?

	Error: 'sqrt' not yet implemented or not found in the registry. This kernel may not be supported by the tfjs backend you have chosen
	    at notYetImplemented (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:169598:11)
	    at MathBackendWebGL.KernelBackend.sqrt (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:169335:16)
	    at /tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:188067:27
	    at /tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172500:55
	    at /tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172333:22
	    at Engine.scopedRun (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172343:23)
	    at Engine.tidy (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172332:21)
	    at kernelFunc (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172500:29)
	    at /tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172512:27
	    at Engine.scopedRun (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172343:23)
	    at Engine.runKernelFunc (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:172510:14)
	    at sqrt_ (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:188066:19)
	    at sqrt__op (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:173278:29)
	    at Tensor.sqrt (/tmp/karma-typescript-bundle--40-PfshppMcwGvg-.js:201056:12)
	    at UserContext.<anonymous> (src/backend_webgl_test.ts:271:22 <- src/backend_webgl_test.js:283:26)
	    at <Jasmine>

[tafsiri]

@mattsoulanille Could you dig into that a bit more? For examples which test is this in, backend_webgl_test.ts calls a number of different ops (e.g. square or relu), which should dispatch to modular kernels. We should find out why those work and sqrt doesnt?

One thing i suspect is that the a new backend is being registered but the kernels are not being copied over into it. Look for uses of tf.copyRegisteredKernels('webgl', ...);, that is probably not being called for the custom backend being used in that test.

Part of the reason i'm asking is that your todo says to remove this once the backend is modularized, but that is what you are doing in this PR, so it is unclear what the blocker is for removing this.

[mattsoulanille]

tf.copyRegisteredKernels fixed it. Thanks!

[tafsiri]

Reviewed 1 of 3 files at r1, 2 of 2 files at r2, 2 of 2 files at r3.
Reviewable status: complete! 1 of 1 approvals obtained