[layers] Add LayerNormalization

[caisq]

• LayerNormalization is used in transformer and transformer-based model topologies

Fixes #1548

FEATURE

---

This change is 

[dsmilkov]

Reviewed 1 of 3 files at r1.
Reviewable status: 0 of 1 approvals obtained (waiting on @caisq, @dsmilkov, and @nsthorat)

---

tfjs-layers/src/exports_layers.ts, line 852 at r1 (raw file):

 * References:
 *   - [Layer Normalization](https://arxiv.org/abs/1607.06450)
 */

missing @doc annotation so it shows up in our API docs

---

tfjs-layers/src/layers/normalization.ts, line 477 at r1 (raw file):

  static className = 'BatchNormalization';

  protected axis: number|number[];

why protected instead of private? do you extend somewhere from LayerNormalization?

---

tfjs-layers/src/layers/normalization.ts, line 524 at r1 (raw file):

    this.supportsMasking = true;
    // TODO(cais): Test masking support.

Why not test in this PR?

---

tfjs-layers/src/layers/normalization.ts, line 597 at r1 (raw file):

      let offset = broadcast(this.beta.read());

      // TODO(cais): The tiling below is a workaround for the limitation of

TODO(issue_link) instead of TODO(username) might be better.

---

tfjs-layers/src/layers/normalization_test.ts, line 905 at r1 (raw file):

      0.1408553570508957
    ]);
  });

can you add unit tests to assert there is no memory leak

[caisq]

Thanks for the review, @dsmilkov

Reviewable status: 0 of 1 approvals obtained (waiting on @dsmilkov and @nsthorat)

---

tfjs-layers/src/exports_layers.ts, line 852 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

missing @doc annotation so it shows up in our API docs

Done.

---

tfjs-layers/src/layers/normalization.ts, line 477 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

why protected instead of private? do you extend somewhere from LayerNormalization?

Changed to private. Same for gamma and beta below.

---

tfjs-layers/src/layers/normalization.ts, line 524 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

Why not test in this PR?

Done.

---

tfjs-layers/src/layers/normalization.ts, line 597 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

TODO(issue_link) instead of TODO(username) might be better.

Done.

---

tfjs-layers/src/layers/normalization_test.ts, line 905 at r1 (raw file):

Previously, dsmilkov (Daniel Smilkov) wrote…

can you add unit tests to assert there is no memory leak

Done. Added leak tests for both forward and training.

[dsmilkov]

Perfect. Thank you!!

Reviewed 1 of 3 files at r1, 3 of 3 files at r2.
Reviewable status: complete! 1 of 1 approvals obtained (waiting on @nsthorat)