@stellarapp/tfjs-stellar 1.0.0

package/src/layers/transformer_decoder.ts

@@ -0,0 +1,236 @@
+import * as tf from "@tensorflow/tfjs";
+import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
+import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
+
+import { type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
+import { CachedRoPEMultiHeadAttention } from "@/layers/cached_rope_multihead_attention";
+
+
+export interface TransformerDecoderArgs extends Omit<MultiHeadAttentionArgs, "causal"> {
+    activation?: "relu" | "gelu";
+    dimsFeedForward?: number;
+    causal?: boolean; // use causal mask for attention on inputs
+}
+
+
+/**
+ * This class implements the transformer decoder architecture from
+ * the 2017 paper "Attention Is All You Need".
+ * 
+ * This decoder-only transformer layer accepts one tensor input.
+ * The input tensor should have the shape
+ * `[ batch, sequences, embedding dims ]`.
+ * 
+ * Causal masking is enabled by default for the initial attention sub-layer.
+ * 
+ * @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 on inputs (masks future inputs to prevent looking ahead), default `true`
+ * @param dropout use dropout during the attention calculations, default `0.1`
+ * @param activation the activation of the intermediate feed forward layer, default `relu`
+ * @param dimsFeedForward the size of the intermediate feed forward layer, default `2048`
+ * @param useBias use bias for the dense sublayers and multiHead attention's dense sublayers, default `true`
+ */
+export class TransformerDecoder extends tf.layers.Layer {
+    static className = "TransformerDecoder";
+
+    protected readonly causalSelfAttention: tf.layers.Layer;
+    protected readonly causalSelfAttentionDropout: tf.layers.Layer;
+    protected readonly causalSelfAttentionNorm: tf.layers.Layer;
+
+    protected readonly feedforward1: tf.layers.Layer;
+    protected readonly feedforward2: tf.layers.Layer;
+    protected readonly feedForwardDropout: tf.layers.Layer;
+    protected readonly feedFowardNorm: tf.layers.Layer;
+
+    protected readonly numHeads: number;
+    protected readonly embedDim: number;
+    protected readonly useBias: boolean;
+    protected readonly dropout: number;
+    protected readonly activation: ActivationIdentifier;
+    protected readonly dimsFeedForward: number;
+
+    constructor({ numHeads, embedDim, useBias, dropout, activation, dimsFeedForward, ...args }: TransformerDecoderArgs) {
+        super(args);
+
+        this.numHeads = numHeads;
+        this.embedDim = embedDim;
+        this.useBias = useBias ?? true;
+        this.dropout = dropout ?? 0.1;
+        this.activation = activation ?? "relu";
+
+        if (this.dropout >= 1) {
+            throw Error(`${this.getClassName()}::constructor dropout must be within [0, 1)`);
+        }
+
+        // in the paper section 3.3, d_model=512 (embedDim) and first dense layer outputs d_ff=2048
+        this.dimsFeedForward = dimsFeedForward ?? embedDim * 4;
+
+        // self attention sub-block
+        this.causalSelfAttention = new CachedRoPEMultiHeadAttention({
+            numHeads: this.numHeads, embedDim: this.embedDim,
+            useBias: this.useBias, dropout: this.dropout,
+            causal: true
+        });
+        this.causalSelfAttentionDropout = tf.layers.dropout({ rate: this.dropout })
+        this.causalSelfAttentionNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
+
+        // feed forward sub-block
+        this.feedforward1 = tf.layers.dense({
+            units: this.dimsFeedForward,
+            activation: this.activation,
+            useBias: this.useBias,
+        });
+        this.feedforward2 = tf.layers.dense({
+            units: this.embedDim,
+            activation: "linear",
+            useBias: this.useBias
+        });
+        this.feedForwardDropout = tf.layers.dropout({ rate: this.dropout });
+        this.feedFowardNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
+    }
+
+
+    /**
+     * Forward propagation
+     * 
+     * @param inputs input tensor
+     * @return the output tensor
+     */
+    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
+        // validate the input tensors
+        if (Array.isArray(inputs) && inputs.length != 1 && inputs.length != 2) {
+            throw Error(`${this.getClassName()}::call ${this.name} expects one input tensor, got ${inputs.length} inputs.`);
+        }
+
+        if (Array.isArray(inputs)) {
+            inputs = inputs[0] as tf.Tensor;
+        }
+
+        // perform forward propagation
+        return tf.tidy(() => {
+            let output = this.causalSelfAttentionBlock(inputs, kwargs);
+            output = this.feedForwardBlock(output, kwargs);
+
+            return output;
+        });
+    }
+
+
+    protected causalSelfAttentionBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+        return tf.tidy(() => {
+            const residual = x;
+
+            let attention = this.causalSelfAttention.apply(x, kwargs) as tf.Tensor;
+            attention = this.causalSelfAttentionDropout.apply(attention, kwargs) as tf.Tensor;
+            attention = tf.add(attention, residual);
+            attention = this.causalSelfAttentionNorm.apply(attention, kwargs) as tf.Tensor;
+
+            return attention;
+        });
+    }
+
+
+    protected feedForwardBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+        return tf.tidy(() => {
+            const residual = x;
+
+            let feedForward = this.feedforward1.apply(x, kwargs);
+            feedForward = this.feedforward2.apply(feedForward, kwargs);
+            feedForward = this.feedForwardDropout.apply(feedForward, kwargs) as tf.Tensor;
+            feedForward = tf.add(feedForward, residual);
+            feedForward = this.feedFowardNorm.apply(feedForward, kwargs) as tf.Tensor;
+
+            return feedForward;
+        });
+    }
+
+
+    /**
+     * Initialize the sublayers' weights and track them to enable serialization
+     */
+    override build(inputShape: tf.Shape | tf.Shape[]): void {
+        let input_shapes: tf.Shape[] = [];
+
+        if (Array.isArray(inputShape) && Array.isArray(inputShape[0])) {
+            // input is an array of shapes
+            input_shapes = inputShape as tf.Shape[];
+        } else if (inputShape.length != 0) {
+            // input is a single shape
+            input_shapes = [inputShape as tf.Shape];
+        }
+
+        if (input_shapes.length != 1 && input_shapes.length != 2) {
+            throw Error(`${this.getClassName()}::build ${this.name} expects an input shape` +
+                ` of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`)
+        }
+
+        const [decoderInputShape] = input_shapes;
+
+        if (decoderInputShape?.length != 3) {
+            throw Error(`${this.getClassName()}::build ${this.name} expects an input shape` +
+                ` of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`)
+        }
+
+        // initialize causal self attention sub-block's weights
+        this.causalSelfAttention.build(decoderInputShape);
+        this.causalSelfAttentionNorm.build(this.causalSelfAttention.computeOutputShape(decoderInputShape));
+
+        // initialize feedforward sub-block's weights
+        const feedforward1OutputShape = this.feedforward1.computeOutputShape(decoderInputShape);
+        const feedforward2OutputShape = this.feedforward2.computeOutputShape(feedforward1OutputShape);
+
+        this.feedforward1.build(decoderInputShape);
+        this.feedforward2.build(feedforward1OutputShape);
+        this.feedFowardNorm.build(feedforward2OutputShape);
+
+        // track sublayers' weights
+        this.trainableWeights = [
+            ...this.causalSelfAttention.trainableWeights,
+            ...this.causalSelfAttentionDropout.trainableWeights,
+            ...this.causalSelfAttentionNorm.trainableWeights,
+            ...this.feedforward1.trainableWeights,
+            ...this.feedforward2.trainableWeights,
+            ...this.feedForwardDropout.trainableWeights,
+            ...this.feedFowardNorm.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);
+    }
+
+
+    /**
+     * Save the layer's hyperparameters for serialization
+     */
+    override getConfig() {
+        const base_config = super.getConfig();
+
+        const config = {
+            numHeads: this.numHeads,
+            embedDim: this.embedDim,
+            useBias: this.useBias,
+            dropout: this.dropout,
+            activation: this.activation,
+            dimsFeedForward: this.dimsFeedForward
+        }
+
+        Object.assign(config, base_config);
+
+        return config;
+    }
+
+}
+
+
+tf.serialization.registerClass(TransformerDecoder);

package/src/layers/transformer_encoder.test.ts

@@ -0,0 +1,85 @@
+import * as tf from "@tensorflow/tfjs";
+
+import { TransformerEncoder } from "@/layers/transformer_encoder";
+
+// disables warning for using the faster node backend,
+// https://github.com/tensorflow/tfjs/issues/5349#issuecomment-885170504
+tf.env().set('IS_NODE', false);
+
+
+describe("TransformerEncoder tests", () => {
+    it("should return an output with the same shape as the input", () => {
+        const input = tf.randomUniform([2, 3, 10]);
+
+        const decoder = new TransformerEncoder({
+            numHeads: 2, embedDim: input.shape.at(-1)!,
+            dropout: 0.5, activation: "gelu", dimsFeedForward: 512, useBias: true
+        });
+
+        const output = decoder.apply(input) as tf.Tensor;
+
+        expect(output.shape.length).toBe(input.shape.length);
+    })
+
+
+    test("correct forward calls", () => {
+        const input = tf.randomUniform([2, 3, 10]);
+
+        const encoder = new TransformerEncoder({ numHeads: 2, embedDim: input.shape.at(-1)! });
+        expect(() => encoder.apply(input)).not.toThrow();
+        expect(() => encoder.apply([input])).not.toThrow();
+
+        const causal = new TransformerEncoder({ numHeads: 2, embedDim: input.shape.at(-1)!, causal: true });
+        expect(() => causal.apply(input)).not.toThrow();
+        expect(() => causal.apply([input])).not.toThrow();
+    })
+
+
+    it("should fail to instantiate a layer if heads count is not divisible by the input's embedding dimension", () => {
+        const input = tf.randomUniform([2, 3, 10]);
+
+        expect(() => new TransformerEncoder({ numHeads: 3, embedDim: input.shape.at(-1)! })).toThrow();
+        expect(() => new TransformerEncoder({ numHeads: 5, embedDim: input.shape.at(-1)! })).not.toThrow();
+    })
+
+
+    it("should not accept non-rank 3 tensor inputs", () => {
+        const incorrect_input = tf.randomUniform([2, 3, 10, 10]);
+        const incorrect_input2 = tf.randomUniform([2, 3]);
+        const correct_input = tf.randomUniform([2, 3, 10]);
+
+
+        const encoder = new TransformerEncoder({ numHeads: 2, embedDim: incorrect_input.shape.at(-1)! });
+        expect(() => encoder.apply([correct_input, correct_input])).toThrow();
+
+        expect(() => encoder.apply(incorrect_input)).toThrow();
+        expect(() => encoder.apply(incorrect_input2)).toThrow();
+
+        expect(() => encoder.apply([correct_input, incorrect_input])).toThrow();
+        expect(() => encoder.apply([incorrect_input, correct_input])).toThrow();
+
+        expect(() => encoder.apply([correct_input, incorrect_input2])).toThrow();
+        expect(() => encoder.apply([incorrect_input2, correct_input])).toThrow();
+    })
+
+
+    it("should accept exactly one input", () => {
+        const input = tf.randomUniform([2, 3, 10]);
+
+        const encoder = new TransformerEncoder({ numHeads: 1, embedDim: input.shape.at(-1)! });
+        expect(() => encoder.apply(input)).not.toThrow();
+        expect(() => encoder.apply([input])).not.toThrow();
+
+        expect(() => encoder.apply([])).toThrow();
+        expect(() => encoder.apply([input, input])).toThrow();
+        expect(() => encoder.apply([input, input, input])).toThrow()
+    })
+
+
+    it("should return a non-empty config dict", () => {
+        const input = tf.randomUniform([2, 3, 10]);
+
+        const encoder = new TransformerEncoder({ numHeads: 1, embedDim: input.shape.at(-1)! });
+        expect(Object.keys(encoder.getConfig())).not.toBe(0);
+    })
+})

package/src/layers/transformer_encoder.ts

@@ -0,0 +1,224 @@
+import * as tf from "@tensorflow/tfjs";
+import { type Kwargs } from "@tensorflow/tfjs-layers/dist/types";
+import { type ActivationIdentifier } from "@tensorflow/tfjs-layers/dist/keras_format/activation_config";
+
+import { MultiHeadAttention, type MultiHeadAttentionArgs } from "@/layers/multihead_attention";
+
+
+export interface TransformerEncoderArgs extends MultiHeadAttentionArgs {
+    activation?: "relu" | "gelu";
+    dimsFeedForward?: number;
+}
+
+
+/**
+ * This class implements the transformer encoder architecture from the 2017 paper
+ * Attention Is All You Need.
+ * 
+ * This layer accepts exactly one tensor input with the shape
+ * `[ batch, sequences, embedding dims ]`.
+ * 
+ * @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` for encoders
+ * @param dropout use dropout during the attention calculations, default `0.1`
+ * @param activation the activation of the intermediate feed forward layer, default `relu`
+ * @param dimsFeedForward the size of the intermediate feed forward layer, default `2048`
+ * @param useBias use bias for the dense sublayers and multiHead attention's dense sublayers, default `true`
+ */
+export class TransformerEncoder extends tf.layers.Layer {
+    static className = "TransformerEncoder";
+
+    private readonly selfAttention: tf.layers.Layer;
+    private readonly selfAttentionDropout: tf.layers.Layer;
+    private readonly selfAttentionNorm: tf.layers.Layer;
+
+    private readonly reluLayer: tf.layers.Layer;
+    private readonly linearLayer: tf.layers.Layer;
+    private readonly feedForwardDropout: tf.layers.Layer;
+    private readonly feedFowardNorm: tf.layers.Layer;
+
+    private readonly numHeads: number;
+    private readonly embedDim: number;
+    private readonly causal: boolean;
+    private readonly useBias: boolean;
+    private readonly dropout: number;
+    private readonly activation: ActivationIdentifier;
+    private readonly dimsFeedForward: number;
+
+
+    constructor({ numHeads, embedDim, causal, useBias, dropout, activation, dimsFeedForward, ...args }: TransformerEncoderArgs) {
+        super(args);
+
+        this.numHeads = numHeads;
+        this.embedDim = embedDim;
+        this.causal = causal ?? false;
+        this.useBias = useBias ?? true;
+        this.dropout = dropout ?? 0.1;
+        this.activation = activation ?? "relu";
+        this.dimsFeedForward = dimsFeedForward ?? 2048;
+
+        if (this.dropout >= 1) {
+            throw Error(`${this.getClassName()}::constructor dropout must be within [0, 1)`);
+        }
+
+        // self attention sub-block
+        this.selfAttention = new MultiHeadAttention({
+            numHeads: this.numHeads, embedDim: this.embedDim, useBias: this.useBias,
+            dropout: this.dropout, causal: this.causal
+        });
+        this.selfAttentionDropout = tf.layers.dropout({ rate: this.dropout })
+        this.selfAttentionNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
+
+        // feed forward sub-block
+        this.reluLayer = tf.layers.dense({
+            units: this.dimsFeedForward, activation: this.activation,
+            useBias: this.useBias
+        });
+        this.linearLayer = tf.layers.dense({
+            units: this.embedDim, activation: "linear",
+            useBias: this.useBias
+        });
+        this.feedForwardDropout = tf.layers.dropout({ rate: this.dropout });
+        this.feedFowardNorm = tf.layers.layerNormalization({ epsilon: 1e-6 });
+    }
+
+
+    /**
+     * Forward propagation
+     */
+    override call(inputs: tf.Tensor | tf.Tensor[], kwargs: Kwargs): tf.Tensor | tf.Tensor[] {
+        // validate the input tensors
+        let input: tf.Tensor;
+
+        if (Array.isArray(inputs)) {
+            if (inputs.length != 1) {
+                throw Error(`${this.getClassName}::call ${this.name} expects exactly 1 tensor` +
+                    ` input, got ${inputs.length} inputs instead.`);
+            }
+
+            input = inputs[0];
+        } else {
+            input = inputs;
+        }
+
+        // perform forward propagation
+        return tf.tidy(() => {
+            const attention = this.selfAttentionBlock(input, kwargs);
+            const feedforward = this.feedForwardBlock(attention, kwargs);
+
+            return feedforward;
+        });
+    }
+
+
+    private selfAttentionBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+        return tf.tidy(() => {
+            const residual = x;
+
+            let attention = this.selfAttention.apply(x, kwargs) as tf.Tensor;
+            attention = this.selfAttentionDropout.apply(attention, kwargs) as tf.Tensor;
+            attention = tf.add(attention, residual);
+            attention = this.selfAttentionNorm.apply(attention) as tf.Tensor;
+
+            return attention;
+        });
+    }
+
+
+    private feedForwardBlock(x: tf.Tensor, kwargs: Kwargs): tf.Tensor {
+        return tf.tidy(() => {
+            const residual = x;
+
+            let feedForward = this.reluLayer.apply(x);
+            feedForward = this.linearLayer.apply(feedForward);
+            feedForward = this.feedForwardDropout.apply(feedForward, kwargs) as tf.Tensor;
+            feedForward = tf.add(feedForward, residual);
+            feedForward = this.feedFowardNorm.apply(feedForward) as tf.Tensor;
+
+            return feedForward;
+        });
+    }
+
+
+    /**
+     * Initialize the sublayers' weights and track them to enable backpropagation.
+     */
+    override build(inputShape: tf.Shape | tf.Shape[]): void {
+        let input_shapes: tf.Shape[] = [];
+
+        if (Array.isArray(inputShape) && Array.isArray(inputShape[0])) {
+            // input is an array of shapes
+            input_shapes = inputShape as tf.Shape[];
+        } else if (inputShape.length != 0) {
+            // input is a single shape
+            input_shapes = [inputShape as tf.Shape];
+        }
+
+        // expects only 1 rank 3 tensor input
+        if (input_shapes.length != 1 || input_shapes[0].length != 3) {
+            throw Error(`${this.getClassName()}::build ${this.name} expects a single input shape of [batch, seq, embed_dim], got ${JSON.stringify(inputShape)}`)
+        }
+
+        // initialize self attention sub-block's weights
+        this.selfAttention.build(inputShape);
+        this.selfAttentionNorm.build(inputShape);
+
+        // inintialize feedforward sub-block's weights
+        const reluLayerOutputShape = this.reluLayer.computeOutputShape(inputShape);
+        const linearLayerOutputShape = this.linearLayer.computeOutputShape(reluLayerOutputShape);
+
+        this.reluLayer.build(inputShape);
+        this.linearLayer.build(reluLayerOutputShape);
+        this.feedFowardNorm.build(linearLayerOutputShape);
+
+        // track sublayers' weights
+        this.trainableWeights = [
+            ...this.selfAttention.trainableWeights,
+            ...this.selfAttentionDropout.trainableWeights,
+            ...this.selfAttentionNorm.trainableWeights,
+            ...this.reluLayer.trainableWeights,
+            ...this.linearLayer.trainableWeights,
+            ...this.feedForwardDropout.trainableWeights,
+            ...this.feedFowardNorm.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);
+    }
+
+
+    /**
+     * Save the layer's hyperparameters for serialization
+     */
+    override getConfig(): tf.serialization.ConfigDict {
+        const base_config = super.getConfig();
+
+        const config = {
+            numHeads: this.numHeads,
+            embedDim: this.embedDim,
+            causal: this.causal,
+            useBias: this.useBias,
+            dropout: this.dropout,
+            activation: this.activation,
+            dimsFeedForward: this.dimsFeedForward
+        };
+
+        Object.assign(config, base_config);
+
+        return config;
+    }
+}
+
+
+tf.serialization.registerClass(TransformerEncoder);

package/src/losses/dice.ts

@@ -0,0 +1,156 @@
+import * as tf from "@tensorflow/tfjs";
+import { categoricalCrossentropy, binaryCrossentropy } from "@tensorflow/tfjs-layers/dist/losses";
+
+const epsilon = 1e-7;
+
+const REDUCE_HW = [1, 2]; // reduce over width and height
+const REDUCE_BHW = [0, 1, 2]; // reduce over batch, width, height
+const REDUCE_BHWC = [0, 1, 2, 3]; // reduce all dimensions
+
+
+// Standard (Sorensen) Dice Loss
+export function diceBinaryStandard(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+
+    const y_true_flat = tf.reshape(y_true, [y_true.shape[0], -1]);
+    const y_pred_flat = tf.reshape(y_pred, [y_pred.shape[0], -1]);
+
+    const intersection = tf.sum(tf.mul(y_true_flat, y_pred_flat), 1);
+    const union = tf.add(tf.sum(y_true_flat, 1), tf.sum(y_pred_flat, 1));
+
+    const dice = tf.div(
+        intersection.mul(2).add(epsilon),
+        union.add(epsilon)
+    );
+
+    return tf.scalar(1).sub(dice);
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceBinaryStandard, "name", { value: "diceBinaryStandard", configurable: false });
+
+
+// https://github.com/keras-team/keras/blob/v3.3.3/keras/src/losses/losses.py#L1983-L2010
+export function diceBinaryGlobal(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+    const y_true_flat = tf.reshape(y_true, [-1]);
+    const y_pred_flat = tf.reshape(y_pred, [-1]);
+
+    const intersection = tf.sum(tf.mul(y_true_flat, y_pred_flat));
+    const union = tf.add(tf.sum(y_true_flat), tf.sum(y_pred_flat));
+
+    const dice = tf.div(
+        intersection.mul(2).add(epsilon),
+        union.add(epsilon)
+    );
+
+    return tf.scalar(1).sub(dice);
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceBinaryGlobal, "name", { value: "diceBinaryGlobal", configurable: false });
+
+
+export function diceCategoricalStandard(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+    const intersection = tf.sum(tf.mul(y_true, y_pred), REDUCE_HW);
+    const union = tf.add(y_true, y_pred).sum(REDUCE_HW);
+
+    const dice = tf.div(
+        intersection.mul(2).add(epsilon),
+        union.add(epsilon)
+    );
+
+    return tf.scalar(1).sub(tf.mean(dice, -1));
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceCategoricalStandard, "name", { value: "diceCategoricalStandard", configurable: false });
+
+
+export function diceCategoricalGeneralized(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+
+    // this is done twice so we calculate it once
+    const y_true_sum = y_true.sum(REDUCE_BHW);
+
+    const weighting = tf.div(1, y_true_sum.square().add(epsilon));
+
+    const intersection = tf.sum(tf.mul(y_true, y_pred), REDUCE_BHW).mul(weighting).sum();
+    const union = tf.add(y_true_sum, y_pred.sum(REDUCE_BHW)).mul(weighting).sum();
+
+    const dice = tf.div(
+        intersection.mul(2).add(epsilon),
+        union.add(epsilon)
+    );
+
+    return tf.scalar(1).sub(dice);
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceCategoricalGeneralized, "name", { value: "diceCategoricalGeneralized", configurable: false });
+
+
+export function diceCategoricalGlobal(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+
+    const intersection = tf.sum(tf.mul(y_true, y_pred), REDUCE_BHWC);
+    const union = tf.add(tf.sum(y_true, REDUCE_BHWC), tf.sum(y_pred, REDUCE_BHWC));
+
+    const dice = tf.div(
+        intersection.mul(2).add(epsilon),
+        union.add(epsilon)
+    );
+
+    return tf.scalar(1).sub(dice);
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceCategoricalGlobal, "name", { value: "diceCategoricalGlobal", configurable: false });
+
+
+/**
+ * Calculates the Sorensen-Dice coefficient and the binary cross entropy losses.
+ * Both have equal weight.
+ * 
+ * @param y_true the label tensor
+ * @param y_pred the prediction tensor (not sparse)
+ * @returns a tensor of shape `[ batch ]`
+ */
+export function diceBinaryCrossentropy(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+    // reduce cross entropy shape from [B, H, W] to [B] to match dice
+    const bce = binaryCrossentropy(y_true, y_pred).mean(REDUCE_HW);
+    const dice = diceBinaryStandard(y_true, y_pred);
+
+    return tf.add(bce.mul(0.5), dice.mul(0.5));
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceBinaryCrossentropy, "name", { value: "diceBinaryCrossentropy", configurable: false });
+
+
+/**
+ * Calculates the Sorensen-Dice coefficient and the categorical cross entropy losses.
+ * Both have equal weight. Expects dense (non-sparse) label tensors.
+ * 
+ * This does not support sparse tensors because TFJS's
+ * sparseCategoricalCrossentropy loss onehots the label
+ * and calls categoricalCrossentropy. See
+ * https://github.com/tensorflow/tfjs/blob/0fc04d958ea592f3b8db79a8b3b497b5c8904097/tfjs-layers/src/losses.ts#L143-L146
+ * 
+ * @param y_true the label  
+ * @param y_pred the prediction tensor (not sparse)
+ * @returns a tensor of shape `[ batch ]`
+ */
+export function diceCategoricalCrossentropy(y_true: tf.Tensor, y_pred: tf.Tensor): tf.Tensor {
+    // reduce cross entropy shape from [B, H, W] to [B] to match dice
+    const cce = categoricalCrossentropy(y_true, y_pred).mean(REDUCE_HW);
+    const dice = diceCategoricalStandard(y_true, y_pred);
+
+    return tf.add(cce.mul(0.5), dice.mul(0.5));
+}
+
+
+// prevents minification of function name which TFJS relies on
+Object.defineProperty(diceCategoricalCrossentropy, "name", { value: "diceCategoricalCrossentropy", configurable: false });