@stellarapp/tfjs-stellar 1.0.0

package/src/layers/multihead_attention.ts

@@ -0,0 +1,371 @@
+import * as tf from '@tensorflow/tfjs';
+import { type LayerArgs } from '@tensorflow/tfjs-layers/dist/engine/topology';
+import { type Kwargs } from '@tensorflow/tfjs-layers/dist/types';
+import { generateCausalAttentionMask } from "@/utils";
+
+
+export interface MultiHeadAttentionArgs extends LayerArgs {
+    numHeads: number;
+    embedDim: number;
+    useBias?: boolean;
+    dropout?: number;
+    causal?: boolean;
+}
+
+
+export interface ScaledDotProductionAttentionKwargs {
+    training?: boolean;
+    dropout?: number;
+    causal?: boolean;
+    scaling_factor?: number;
+}
+
+
+/**
+ * This MultiHead Attention layer implements the algorithm as described in
+ * the paper "Attention is all you Need" Vaswani et al., 2017.
+ * 
+ * @param numHeads number of attention heads to use
+ * @param embedDim the embedding size of the input (input embeddings, typically the last dimension)
+ * @param causal use causal masking, default `false`
+ * @param dropout use dropout during the attention calculations, default `0.0`
+ * @param useBias use bias for the dense sublayers, default `true`
+ * 
+ * The TensorFlow version uses tf.einsum, whose gradient op has not yet been
+ * implemented (https://github.com/tensorflow/tfjs/pull/4955#discussion_r619219334),
+ * therefore we follow the PyTorch implementation described in:
+ * https://docs.pytorch.org/tutorials/intermediate/transformer_building_blocks.html#multiheadattention
+ * https://docs.pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html
+ * 
+ * This implementation is different from TensorFlow's whose attention weights
+ * are shaped [embed dim, heads, embed dim] where as PyTorch and OpenAI's attention weights
+ * are shaped [embed dim, embed dim]
+ * https://github.com/pytorch/pytorch/blob/134179474539648ba7dee1317959529fbd0e7f89/torch/nn/modules/activation.py#L1080
+ * https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/model.py#L53
+ * 
+ * TODO: implement a fast track for self attention (query = key = value)
+ * where a single dense layer combines and replaces the query, key and projection layers
+ * 
+ * TODO: add kDim and vDim to accept key and values whose embedding dimensions differ from query's.
+ */
+export class MultiHeadAttention extends tf.layers.Layer {
+    static className = "MultiHeadAttention";
+    protected readonly numHeads: number;
+    protected readonly embedDim: number; // size of embedding dim of inputs, also per attention head
+    protected readonly useBias: boolean;
+    protected readonly dropout: number;
+    protected readonly causal: boolean; // use causal attention to mask future words
+
+    // projection simply means matrix multiplying query, key, and value
+    // with weights to create a representation of the inputs
+    protected readonly queryProjection: tf.layers.Layer;
+    protected readonly keyProjection: tf.layers.Layer;
+    protected readonly valueProjection: tf.layers.Layer;
+    protected readonly outputProjection: tf.layers.Layer;
+
+
+    constructor({ numHeads, embedDim, useBias = true, dropout = 0.0, causal = false, ...args }: MultiHeadAttentionArgs) {
+        super(args);
+
+        if (embedDim % numHeads != 0) {
+            throw Error(`${this.getClassName()}::constructor ${this.name} embedDim (${embedDim}) is not divisible by numHeads (${numHeads})`);
+        }
+
+        this.numHeads = numHeads;
+        this.embedDim = embedDim;
+        this.useBias = useBias;
+        this.dropout = dropout;
+        this.causal = causal;
+
+        if (this.dropout >= 1) {
+            throw Error(`${this.getClassName()}::constructor dropout must be within [0, 1)`);
+        }
+
+        // intialize the projection weights, this should be in the
+        // build() function but is done here to avoid linting complaints
+        this.queryProjection = tf.layers.dense({ useBias, units: embedDim });
+        this.keyProjection = tf.layers.dense({ useBias, units: embedDim });
+        this.valueProjection = tf.layers.dense({ useBias, units: embedDim });
+        this.outputProjection = tf.layers.dense({ useBias, units: embedDim });
+    }
+
+
+    /**
+     * Forward propagation. Provide one input tensor or three identical tensors to self-attention.
+     * @param inputs a single tensor for self-attention or an array of exactly three
+     *   tensors that are either identical (self-attention) or different (cross-attention)
+     * @param kwargs.packingMask a mask to prevent tokens from attending across document boundaries
+     */
+    override call(
+        inputs: tf.Tensor | tf.Tensor[],
+        kwargs: Kwargs & {
+            packingMask?: tf.Tensor,
+            causalMask?: tf.Tensor,
+        }
+    ): tf.Tensor | tf.Tensor[] {
+        // validate the input tensors
+        if (!Array.isArray(inputs)) {
+            inputs = [inputs];
+        }
+
+        // accept only 1 input (self attention) or 3 inputs (self or cross attention)
+        if (inputs.length != 1 && inputs.length != 3) {
+            throw Error(`${this.getClassName()}::call ${this.name} expects exactly one or three input tensors, ${inputs.length} were provided`);
+        }
+
+        for (const input of inputs) {
+            if (input.shape.length != 3) {
+                throw Error(`${this.getClassName()}::call ${this.name} expected input shapes of [batch, seq, embed_dim], got ${JSON.stringify(input.shape)}`);
+            }
+        }
+
+        const [query, key, value] = inputs;
+        const packingMask = kwargs.packingMask ?? null;
+        const causalMask = kwargs.causalMask ?? null;
+
+        return inputs.length == 3
+            // cross-attention
+            ? this.forward(query!, key!, value!, packingMask, causalMask, kwargs)
+            // self-attention
+            : this.forward(query!, query!, query!, packingMask, causalMask, kwargs);
+    }
+
+
+    /**
+     * Forward propagation
+     */
+    protected forward(
+        query_input: tf.Tensor,
+        key_input: tf.Tensor,
+        value_input: tf.Tensor,
+        packing_mask: tf.Tensor | null,
+        causal_mask: tf.Tensor | null,
+        kwargs: Kwargs): tf.Tensor {
+
+        // dimensions abbreviations
+        // batch = the number of sequences in the input
+        // seq = the length of each sequence in the input
+        // dims = the size of each token's embedding
+        return tf.tidy(() => {
+            const { query, key, value } = this.applyInputProjections(query_input, key_input, value_input);
+
+            // swap the seq and heads dimensions: [batch, seq, heads, head_dim] -> [batch, heads, seq, head_dim]
+            const move_head_dim_forward = [0, 2, 1, 3];
+
+            const {
+                query_split, key_split, value_split
+            } = this.splitHeads(query, key, value, move_head_dim_forward);
+
+            // apply scaled dot production attention to get [batch, seq, numHeads, embedDim]
+            const spda = MultiHeadAttention.scaledDotProductionAttention(
+                query_split, key_split, value_split,
+                kwargs.attentionMask ?? null, packing_mask, causal_mask,
+                this.dropout, this.causal, kwargs);
+
+            // concat heads and apply the output projection
+            const output = this.outputProjection.apply(
+                spda.transpose(move_head_dim_forward).reshape([query_input.shape[0], -1, this.embedDim]));
+
+            return output as tf.Tensor;
+        })
+    }
+
+
+    protected applyInputProjections(query_input: tf.Tensor, key_input: tf.Tensor, value_input: tf.Tensor) {
+        // apply input projections, this is a batched matrix multiplication operated on the last
+        // dimension of query_input and first dimension of the dense layer weights,
+        // [batch, seq, dims] x [dims, dims] = [batch x seq, dims] x [dims, dims] = [batch x seq, dims] = [batch, seq, dims]
+        return tf.tidy(() => {
+            return {
+                query: this.queryProjection.apply(query_input) as tf.Tensor,
+                key: this.keyProjection.apply(key_input) as tf.Tensor,
+                value: this.valueProjection.apply(value_input) as tf.Tensor
+            }
+        })
+    }
+
+
+    protected splitHeads(query: tf.Tensor, key: tf.Tensor, value: tf.Tensor, shuffle: number[]) {
+        // split heads and prepare for scaled dot product attention by splitting the
+        // last dimension to get the heads, bring the heads forward
+        // [batch, seq, dims] -> [batch, seq, heads, dims / heads] -> [batch, heads, seq, head_dim]
+        const batch_size = query.shape[0];
+        const split_heads = [batch_size, -1, this.numHeads, this.embedDim / this.numHeads];
+
+        return tf.tidy(() => {
+            return {
+                query_split: query.reshape(split_heads).transpose(shuffle) as tf.Tensor4D,
+                key_split: key.reshape(split_heads).transpose(shuffle) as tf.Tensor4D,
+                value_split: value.reshape(split_heads).transpose(shuffle) as tf.Tensor4D
+            }
+        })
+    }
+
+
+    /**
+     * Applies the scaled dot-product formula: softmax(QK_t / sqrt(d_k))V,
+     * formula (1) of the 2017 paper Attention Is All You Need
+     * 
+     * @param attentionMask a mask to prevent tokens from being 
+     *   attended to (usually for padding tokens). It should have the shape
+     *   [batch, head, query_sequence_len, key_sequence_len]. To use in
+     *   conjunction with causal masking, the tensor should be a boolean type
+     *   where false indicates a masked token.
+     * @param packingMask a mask to prevent tokens from attending across document boundaries
+     */
+    static scaledDotProductionAttention(
+        query: tf.Tensor,
+        key: tf.Tensor,
+        value: tf.Tensor,
+        attentionMask: tf.Tensor | null,
+        packingMask: tf.Tensor | null,
+        causalMask: tf.Tensor | null,
+        dropout: number,
+        causal: boolean,
+        kwargs: ScaledDotProductionAttentionKwargs = {}
+    ): tf.Tensor {
+        return tf.tidy(() => {
+            const { training = false, scaling_factor } = kwargs;
+
+            key.shape.forEach((val, index) => {
+                if (key.shape[index] != value.shape[index]) {
+                    throw Error(`scaledDotProductionAttention: expected key and value` +
+                        ` to have the same shape, got ${JSON.stringify(key.shape)} (key) and` +
+                        ` ${JSON.stringify(value.shape)} (value)`);
+                }
+            })
+
+
+            // mask's shape is [..., seq, seq] where seq is the number of words/tokens in the input,
+            // not adding the batch dimension yet to lessen the calculations
+            const causal_mask_shape = [
+                query.shape[query.shape.length - 2],
+                key.shape[key.shape.length - 2]];
+
+            let mask = tf.zeros(causal_mask_shape);
+
+            if (causal && causal_mask_shape[0] > 1) {
+                if (attentionMask && attentionMask.dtype != "bool") {
+                    throw Error(`scaledDotProductionAttention: the attention mask must be undefined or a boolean type if used with causal attention`);
+                }
+
+                // apply a causal attention mask so that tokens can only attend to preceding tokens,
+                // prevents looking at head
+                if (causalMask) {
+                    mask = causalMask;
+                } else {
+                    mask = generateCausalAttentionMask(causal_mask_shape[0], causal_mask_shape[1]);
+                }
+            }
+
+            if (attentionMask) {
+                if (attentionMask.dtype == "bool") {
+                    // convert the boolean mask to float
+                    // warning: do not use 1e9, it will overflow, use something smaller like 1e7
+                    mask = mask.add(attentionMask.cast("float32").sub(1).mul(1e7));
+                } else {
+                    // this will occur only when not using causal masking,
+                    // if the attention mask is not boolean, it's assumed the masking is already calculated,
+                    mask = attentionMask;
+                }
+            }
+
+            // 1. matrix multiply query and transposed key
+            // 2. divide by scaling factor
+            // 3. apply softmax to the result
+            // 4. apply attention and/or causal mask
+            // 5. apply dropout
+            // 6. matrix multiply softmax result with value
+            let pre_softmax = query
+                .matMul(key, false, true)
+                .div(Math.sqrt(scaling_factor ?? key.shape[key.shape.length - 1]))
+                .add(mask);
+
+            if (packingMask) {
+                // packing mask is added separately because each mask within a batch may be different,
+                // so it cannot be broadcasted
+                pre_softmax = pre_softmax.add(packingMask);
+            }
+
+            const spda = tf.softmax(pre_softmax);
+
+            const spda_dropout = tf.dropout(spda, training ? dropout : 0);
+            const attention = spda_dropout.matMul(value);
+
+            return attention;
+        });
+    }
+
+
+    override build(inputShape: tf.Shape | tf.Shape[]): void {
+        let input_shape: tf.Shape[] = [];
+
+        if (Array.isArray(inputShape) && Array.isArray(inputShape[0])) {
+            input_shape = inputShape as tf.Shape[];
+        } else {
+            input_shape = [inputShape as tf.Shape, inputShape as tf.Shape, inputShape as tf.Shape];
+        }
+
+        if (input_shape.length != 1 && input_shape.length != 3) {
+            throw Error(`${this.getClassName()}::build ${this.name} accepts either exactly one or three inputs, received ${JSON.stringify(inputShape)}`);
+        }
+
+        // initialize the sublayer weights
+        this.queryProjection.build(input_shape[0]);
+        this.keyProjection.build(input_shape[1]);
+        this.valueProjection.build(input_shape[2]);
+        this.outputProjection.build(input_shape[0]);
+
+        // the sublayer weights need to be tracked by this layer otherwise
+        // backpropagation will complain about no trainable parameters found,
+        // this is an extra step that TF's Python version does not need
+        this.trainableWeights = [
+            ...this.queryProjection.trainableWeights,
+            ...this.keyProjection.trainableWeights,
+            ...this.valueProjection.trainableWeights,
+            ...this.outputProjection.trainableWeights
+        ];
+
+        // rename the weights otherwise they'll take on the default naming and overlap
+        // each other which breaks model loading due to duplicate weight names
+        let indexing = 0;
+
+        for (const weight of this.trainableWeights) {
+            const unique_name = `${this.getClassName()}_${indexing}`;
+            (weight as any).name += unique_name;
+            (weight as any).originalName += unique_name;
+            indexing++;
+        }
+
+        super.build(inputShape);
+    }
+
+
+    /**
+     * MultiHead attention's output is the same shape the query's.
+     */
+    override computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[] {
+        return Array.isArray(inputShape) && Array.isArray(inputShape[0]) ? inputShape[0] : inputShape;
+    }
+
+
+    override getConfig() {
+        const base_config = super.getConfig();
+
+        const config = {
+            numHeads: this.numHeads,
+            embedDim: this.embedDim,
+            useBias: this.useBias,
+            causal: this.causal,
+            dropout: this.dropout,
+            name: this.name,
+        }
+
+        Object.assign(config, base_config);
+
+        return config;
+    }
+}
+
+
+tf.serialization.registerClass(MultiHeadAttention);

package/src/layers/positional_encoding.test.ts

@@ -0,0 +1,113 @@
+import * as tf from '@tensorflow/tfjs';
+
+import { PositionalEncoding } from '@/layers/positional_encoding';
+
+// disables warning for using the faster node backend,
+// https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
+tf.env().set('IS_NODE', false);
+
+
+describe("PositionalEncoding tests", () => {
+    it("should fail to instantiate a layer", () => {
+        expect(() => new PositionalEncoding({ maxSequenceLength: 3, embedDim: 0 })).toThrow();
+        expect(() => new PositionalEncoding({ maxSequenceLength: 3, embedDim: -1 })).toThrow();
+        expect(() => new PositionalEncoding({ maxSequenceLength: 0, embedDim: 32 })).toThrow();
+        expect(() => new PositionalEncoding({ maxSequenceLength: -1, embedDim: 32 })).toThrow();
+    })
+
+
+    test("successfull forward calls", () => {
+        const embed_dims = 32;
+        const sequences = 4;
+        const input = tf.randomUniform([2, sequences, embed_dims]);
+
+        const positional = new PositionalEncoding({ embedDim: embed_dims });
+        expect(() => positional.apply(input)).not.toThrow();
+        expect(() => positional.apply([input])).not.toThrow();
+        expect(positional.computeOutputShape(input.shape)).toEqual(input.shape);
+    })
+
+
+    it("should throw when input sequences are too large, embedding dims don't match, input aren't rank 3", () => {
+        const sequences_too_long = tf.randomUniform([100, 32]);
+        const embeddings_too_large = tf.randomUniform([32, 100]);
+        const wrong_rank = tf.randomUniform([10, 32, 32]);
+
+        const positional = new PositionalEncoding({ maxSequenceLength: 10, embedDim: 32 });
+
+        expect(() => positional.apply(sequences_too_long)).toThrow();
+        expect(() => positional.apply(embeddings_too_large)).toThrow();
+        expect(() => positional.apply(wrong_rank)).toThrow();
+    })
+
+
+    it("should return a non-empty config dict", () => {
+        const attention = new PositionalEncoding({ embedDim: 32 });
+        expect(Object.keys(attention.getConfig())).not.toBe(0);
+    })
+
+
+    // PyTorch implementation at found at
+    // https://pytorch-tutorials-preview.netlify.app/beginner/transformer_tutorial.html
+    it("should be within 1e-6 of PyTorch's implementation", () => {
+        const pytorch_embed4 = tf.tensor([
+            [[0.0000000, 1.0000000, 0.0000000, 1.0000000],
+            [0.8414710, 0.5403023, 0.0099998, 0.9999500],
+            [0.9092974, -0.4161468, 0.0199987, 0.9998000],
+            [0.1411200, -0.9899925, 0.0299955, 0.9995500],
+            [-0.7568025, -0.6536436, 0.0399893, 0.9992001],
+            [-0.9589243, 0.2836622, 0.0499792, 0.9987503],
+            [-0.2794155, 0.9601703, 0.0599640, 0.9982005],
+            [0.6569866, 0.7539023, 0.0699428, 0.9975510],
+            [0.9893582, -0.1455000, 0.0799147, 0.9968017],
+            [0.4121185, -0.9111302, 0.0898785, 0.9959527]]]);
+
+        const pytorch_embed8 = tf.tensor([
+            [[0.0000000e+00, 1.0000000e+00, 0.0000000e+00, 1.0000000e+00,
+                0.0000000e+00, 1.0000000e+00, 0.0000000e+00, 1.0000000e+00],
+            [8.4147096e-01, 5.4030234e-01, 9.9833414e-02, 9.9500418e-01,
+                9.9998331e-03, 9.9994999e-01, 9.9999981e-04, 9.9999952e-01],
+            [9.0929741e-01, -4.1614684e-01, 1.9866931e-01, 9.8006660e-01,
+                1.9998666e-02, 9.9980003e-01, 1.9999985e-03, 9.9999803e-01],
+            [1.4112000e-01, -9.8999250e-01, 2.9552019e-01, 9.5533651e-01,
+                2.9995499e-02, 9.9955004e-01, 2.9999954e-03, 9.9999553e-01],
+            [-7.5680250e-01, -6.5364361e-01, 3.8941833e-01, 9.2106098e-01,
+                3.9989334e-02, 9.9920011e-01, 3.9999890e-03, 9.9999201e-01],
+            [-9.5892429e-01, 2.8366220e-01, 4.7942552e-01, 8.7758255e-01,
+                4.9979165e-02, 9.9875027e-01, 4.9999789e-03, 9.9998754e-01],
+            [-2.7941549e-01, 9.6017027e-01, 5.6464243e-01, 8.2533562e-01,
+                5.9964005e-02, 9.9820054e-01, 5.9999637e-03, 9.9998200e-01],
+            [6.5698659e-01, 7.5390226e-01, 6.4421761e-01, 7.6484221e-01,
+                6.9942847e-02, 9.9755102e-01, 6.9999420e-03, 9.9997550e-01],
+            [9.8935825e-01, -1.4550003e-01, 7.1735609e-01, 6.9670677e-01,
+                7.9914689e-02, 9.9680167e-01, 7.9999138e-03, 9.9996799e-01],
+            [4.1211849e-01, -9.1113025e-01, 7.8332686e-01, 6.2160999e-01,
+                8.9878544e-02, 9.9595273e-01, 8.9998785e-03, 9.9995953e-01]]]);
+
+        const positional4 = new PositionalEncoding({ embedDim: 4, maxSequenceLength: 10 });
+        positional4.build([]);
+
+        const positional8 = new PositionalEncoding({ embedDim: 8, maxSequenceLength: 10 });
+        positional8.build([]);
+
+        const margin_of_error = 1e-6;
+
+        // the difference between this and PyTorch's implementation
+        //should be insignificantly small
+        expect((positional4.getWeights()[0]
+            .sub(pytorch_embed4)
+            .abs()
+            .arraySync() as [])
+            .flat(2)
+            .filter(i => i > margin_of_error)
+            .length).toBe(0);
+
+        expect((positional8.getWeights()[0]
+            .sub(pytorch_embed8)
+            .abs()
+            .arraySync() as [])
+            .flat(2)
+            .filter(i => i > margin_of_error)
+            .length).toBe(0);
+    });
+});

package/src/layers/positional_encoding.ts

@@ -0,0 +1,158 @@
+import * as tf from "@tensorflow/tfjs";
+import { type LayerArgs } from '@tensorflow/tfjs-layers/dist/engine/topology';
+import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
+
+
+export interface PositionalEncodingArgs extends LayerArgs {
+    // embedding size of each word/token, aka d_model from the paper
+    embedDim: number;
+    // the max length of each sentence, any more or less are truncated or padded
+    maxSequenceLength?: number;
+}
+
+
+/**
+ * This class implements the position encoding logic described in the
+ * 2017 paper "Attention Is All You Need".
+ * 
+ * This layer is untrainable and accepts inputs of shape `[ batch, sequences, embedding dims ]`
+ * and adds positional encoding to return an output tensor of the same shape.
+ * 
+ * @param embedDim the size of each token/word's embedding
+ * @param maxSequenceLength the max number of tokens (words) per input (sentence), default `5120`
+ */
+export class PositionalEncoding extends tf.layers.Layer {
+    static className = "PositionalEncoding";
+    private readonly maxSequenceLength: number;
+    private readonly embedDim: number;
+    private positionalEncodings: tf.LayerVariable;
+
+
+    constructor(args: PositionalEncodingArgs) {
+        super(args);
+
+        this.maxSequenceLength = args.maxSequenceLength ?? 5120;
+        this.embedDim = args.embedDim;
+
+        if (this.maxSequenceLength < 1) {
+            throw Error(`${this.getClassName()}::constructor ${this.name} maxSequenceLength` +
+                ` (${args.maxSequenceLength}) must be greater than 0`);
+        }
+
+        if (this.embedDim < 1) {
+            throw Error(`${this.getClassName()}::constructor ${this.name} embedDim` +
+                ` (${args.embedDim}) must be greater than 0`);
+        }
+
+        // positional encodings are not trainable
+        this.positionalEncodings = this.addWeight('positional_encodings',
+            [this.maxSequenceLength, this.embedDim], "float32",
+            tf.initializers.zeros(), undefined, false);
+    }
+
+
+    /**
+     * Forward propagation. Injects positional encoding to the input embeddings
+     */
+    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
+        // validate the input tensors
+        const input = Array.isArray(inputs) ? inputs[0] : inputs;
+        const sequences = input.shape[1]!;
+
+        if (input.shape.length != 3 || input.shape[2] != this.embedDim) {
+            throw Error(`${this.getClassName()}::call ${this.name} expected an input shape of` +
+                ` [batch, (up to ${this.maxSequenceLength}), ${this.embedDim}], instead got ${input.shape}`);
+        }
+
+        if (sequences > this.maxSequenceLength) {
+            // unexpected sequence length
+            throw Error(`${this.getClassName()}::call ${this.name} received an input with` +
+                ` sequence length (${sequences}) which is greater than the max sequence length` +
+                ` ${this.maxSequenceLength}`);
+        }
+
+        // perform forward propagation
+        return tf.tidy(() => {
+            return input.add(this.positionalEncodings.read()
+                .slice([0, 0], [sequences, this.embedDim]) // gets the first "sequences" rows
+                .expandDims(0)); // introduce the batch dimension and let add() broadcast it
+        })
+    }
+
+    /**
+     * Generate the positional encoding from the paper Attention Is All You Need.
+     * Note that because the inner term of the position formula is the same for both even
+     * and odd indices, we only create half of it and apply sine and cosine individually.
+     */
+    override build(inputShape: tf.Shape | tf.Shape[]): void {
+        tf.tidy(() => {
+            const embedDimHalved = Math.ceil(this.embedDim / 2);
+
+            // create the position matrix as [ 0, 1, 2, 3, etc ],
+            // and broadcast it horizontally to match the number of embeddings,
+            const numerator = tf.range(0, this.maxSequenceLength, 1)
+                .reshape([this.maxSequenceLength, 1])
+                // this creates an extra, unsued positional encoding column later on for odd embedding sizes
+                .broadcastTo([this.maxSequenceLength, embedDimHalved]);
+
+            // the inner term's denominator's exponent's numerator is created as
+            // [ 0, 0, 2, 2, 4, 4, etc ] ( technically [0, 2, 4] as explained above ) and not
+            // [ 0, 2, 4, 6, 8, 10, etc ] because the even and odd indices are counted as pairs
+            // when incrementing "i",
+            // the denominator formula is 10_000^(2i/d_model) where each "i" is a sine cosine pair
+            const denominator = tf.pow(10_000, tf.range(0, this.embedDim, 2).div(this.embedDim));
+
+            const inner_term = numerator.div(denominator);
+
+            const sine = tf.sin(inner_term);
+            const cosine = tf.cos(inner_term);
+
+            // horizontally interweave the sine and cosine columns together to form
+            // [sin, cos, sin, cos, etc]
+            // [sin, cos, sin, cos, etc]
+            // etc
+            const interweaved = [];
+            const ALL_ROWS = -1;
+            const ONE_COL = 1;
+            const FIRST_ROW = 0;
+
+            for (let targetCol = 0; targetCol < this.embedDim / 2; targetCol++) {
+                interweaved.push(sine.slice([FIRST_ROW, targetCol], [ALL_ROWS, ONE_COL]))
+
+                if (targetCol != Math.floor(this.embedDim / 2)) {
+                    // for odd numbered embedDim sizes skip the last cosine column
+                    // e.g. if embedDim = 5, create [ i=0 (sin), i=0 (cos), i=1 (sin), i=1 (cos), i=2 (sin) ]
+                    // and the final i=2 (cos) is ignored
+                    interweaved.push(cosine.slice([FIRST_ROW, targetCol], [ALL_ROWS, ONE_COL]))
+                }
+            }
+
+            // add the positional encoding
+            this.setWeights([tf.concat(interweaved, 1)]);
+        });
+
+        super.build(inputShape);
+    }
+
+
+    override computeOutputShape(inputShape: tf.Shape | tf.Shape[]): tf.Shape | tf.Shape[] {
+        return inputShape;
+    }
+
+
+    override getConfig(): tf.serialization.ConfigDict {
+        const base_config = super.getConfig();
+
+        const config = {
+            maxSequenceLength: this.maxSequenceLength,
+            embedDim: this.embedDim,
+        }
+
+        Object.assign(config, base_config);
+
+        return config;
+    }
+}
+
+
+tf.serialization.registerClass(PositionalEncoding);