The layer includes a custom hashing function. Hash functions are often used in neural networks for dimensionality

idNorm_py.py

import torch
import torch.nn as nn
import math

class IdNorm(nn.Module):
    def __init__(self, cluster_size=128):
        super(IdNorm, self).__init__()
        self.cluster_size = cluster_size
        self.n = None  # Será definido no método build
        self.embs = nn.ModuleList()
        self.random_bias = nn.ParameterList()

    def hash_function(self, vec, n=1):
        a = 0
        b = 0
        mod = (self.cluster_size / 2) - 1
        res = []

        if n == 1:
            for i in vec:
                a = (i + a) % mod
                b = (a + b) % mod
                res.append(round(b))
        else:
            nm1 = n - 1
            for i, val in enumerate(vec):
                a = (val + a) % mod
                b = (a + b) % mod
                if i % n == nm1:
                    res.append(round(b))
        
        return torch.tensor(res, dtype=torch.int64)

    def build(self, input_shape):
        self.n = round(math.log2(input_shape[1]))
        for k in range(self.n):
            self.embs.append(nn.Embedding(self.cluster_size + 1, input_shape[2]))
            self.random_bias.append(nn.Parameter(torch.zeros(1)))

    def forward(self, input):
        if self.n is None:
            self.build(input.shape)

        # Implementar lógica de processamento do método call aqui...
        # Esta parte do código requer adaptação da lógica específica do TensorFlow.js para o PyTorch.

        return input

    def get_config(self):
        return {
            'depth': self.cluster_size,
            'outputDim': self.n
        }

idNorm_tf.py

import tensorflow as tf
import numpy as np

class IdNorm(tf.keras.layers.Layer):
    def __init__(self, cluster_size=128):
        super(IdNorm, self).__init__()
        self.cluster_size = cluster_size

    def hash(self, vec, modt=128, n=1):
        a = 0
        b = 0
        mod = (self.cluster_size / 2) - 1
        res = []

        if n == 1:
            for i in range(len(vec)):
                a = (vec[i] + a) % mod
                b = (a + b) % mod
                res.append(tf.round(b))
        else:
            nm1 = n - 1
            for i in range(len(vec)):
                a = (vec[i] + a) % mod
                b = (a + b) % mod
                if i % n == nm1:
                    res.append(tf.round(b))

        return res

    def build(self, input_shape):
        self.embs = []
        self.n = int(np.round(np.log2(input_shape[1])))
        self.random_bias = []

        for k in range(self.n):
            embedding = tf.keras.layers.Embedding(
                input_dim=self.cluster_size + 1,
                output_dim=input_shape[2],
                name='emb' + str(k)
            )
            self.embs.append(embedding)
            random_bias = self.add_weight(
                name="idnormRandomBias" + str(k),
                shape=(1,),
                initializer=tf.initializers.zeros(),
                dtype=tf.float32,
                trainable=True
            )
            self.random_bias.append(random_bias)

        super(IdNorm, self).build(input_shape)

    def call(self, inputs):
        fr = inputs[0]

        def hash_op(temp, op):
            temp = self.hash(temp, *op)
            return temp

        result = None
        d = inputs[0]
        f = tf.round(tf.multiply(tf.math.tanh(d), 1000))
        f = f.numpy().tolist()
        emb1 = []

        for u in f:
            emb1_u = []
            for u2 in u:
                fl = u2

                MOD_ADLER = (self.cluster_size / 2) - 1
                data = fl
                len_data = len(fl)
                a, b = 1, 0
                index = 0

                while index < len_data:
                    a = (a + data[index]) % MOD_ADLER
                    b = (b + a) % MOD_ADLER
                    index += 1

                emb1_u.append(b)
            emb1.append(emb1_u)

        for u in emb1:
            ids = [tf.round(s + (self.cluster_size / 2)) for s in u]
            ops = [[self.cluster_size, 2] for _ in range(self.n)]
            temp = ids.copy()
            res = [hash_op(temp, op) for op in ops]

            final = []
            for s, i in enumerate(ids):
                tokens = [i]
                for x in range(self.n):
                    tokens.append(res[x][i // (2 ** x)] if i // (2 ** x) < len(res[x]) else 0)
                final.append(tokens.copy())

            for emb in self.embs:
                emb_input = tf.convert_to_tensor([final[s][self.embs.index(emb)] for s in range(len(final))])
                fr = fr + emb(emb_input)

        return fr

idNorm_tfjs.js

tf.layers.Layer.prototype.addSubLayer = function (subLayer, inputShape) {
    subLayer.build(inputShape);
    subLayer.built = true;
    subLayer.weights.forEach(weight => {
        if (!weight.nameModified) {
            weight.nameModified = true;
            const scopedName = weight.name.split('/');
            weight.name = `${this.name}/${subLayer.name}/${scopedName[scopedName.length - 1]}`;
            const scopedOriginalName = weight.originalName.split('/');
            weight.originalName = `${this.name}/${subLayer.name}/${scopedOriginalName[scopedOriginalName.length - 1]}`;
        } else {
            weight.name = `${this.name}/${weight.name}`;
            weight.originalName = `${this.name}/${weight.originalName}`;
        }
        if (this._addedWeightNames.indexOf(weight.name) !== -1) {
            throw new Error(`Duplicate weight name ${weight.name} for layer ${this.name}`);
        }
        this._addedWeightNames.push(weight.name);
    });
    subLayer.trainableWeights.forEach(weight => {
        this._trainableWeights.push(weight);
    });
    subLayer.nonTrainableWeights.forEach(weight => {
        this._nonTrainableWeights.push(weight);
    });
    subLayer.losses.forEach(loss => this.addLoss(loss));
    return subLayer;
};
class idNorm extends tf.layers.Layer {
    static className = 'idNorm';

    constructor(config = {
        clusterSize: 128
    }) {
        super(config);

        this.n = config.n;
        this.clusterSize = config.clusterSize;
    }

    // Hash function
    hash(vec, modt = 128, n = 1) {
        let a = 0;
        let b = 0;
        let mod = (this.clusterSize / 2) - 1;
        let res = [];

        if (n === 1) {
            for (let i = 0; i < vec.length; i++) {
                a = (vec[i] + a) % mod;
                b = (a + b) % mod;
                res.push(Math.round(b));
            }
        } else {
            let nm1 = n - 1;
            for (let i = 0; i < vec.length; i++) {
                a = (vec[i] + a) % mod;
                b = (a + b) % mod;
                if (i % n === nm1) {
                    res.push(Math.round(b));
                }
            }
        }

        return res;
    }

    build(inputShape) {
        this.embs = [];
        this.n = Math.round(Math.log2(inputShape[1]));
        this.randomBias = [];

        for (let k = 0; k < this.n; k++) {
            this.embs[k] = this.addSubLayer(tf.layers.embedding({
                outputDim: inputShape[2],
                inputDim: this.clusterSize + 1,
                name: 'emb' + k
            }), [inputShape[2], opt.length]);

            this.randomBias[k] = this.addWeight("idnormRandomBias" + k, [1], "float32", initializers.zeros([1]), undefined, true);
        }

        this.built = true;
    }

    call(input) {
        let fr = input[0];
        return tf.tidy(() => {
            let result = null;
            let d = input[0];
            let f = d.tanh().mul(1000).round();
            f = f.round().arraySync();
            let emb1 = [];

            for (let u in f) {
                emb1[u] = [];
                for (let u2 in f[u]) {
                    let fl = f[u][u2];
                    let MOD_ADLER = (this.clusterSize / 2) - 1;
                    let data = fl;
                    let len = fl.length;
                    let a = 1,
                        b = 0;
                    let index;

                    for (index = 0; index < len; index++) {
                        a = (a + data[index]) % MOD_ADLER;
                        b = (b + a) % MOD_ADLER;
                    }

                    emb1[u][u2] = b;
                }
            }

            for (let u in emb1) {
                let ids = emb1[u].map(s => Math.round(s + (this.clusterSize / 2)));
                let ops = new Array(this.n).fill(0).map(s => [this.clusterSize, 2]);
                let temp = ids.slice();
                let res = ops.map(s => {
                    let input = [temp].concat(s);
                    temp = this.hash(...input);
                    return temp;
                });

                let final = [];
                ids.map((s, i) => {
                    let tokens = [s];
                    for (let x = 0; x < this.n; x++) {
                        tokens.push(res[x][Math.floor(i / (2 ** x))] || 0);
                    }

                    final.push(tokens.slice());
                });

                for (let emb in this.embs) {
                    fr = tf.add(fr, this.embs[emb].apply(tf.tensor(final.map(s => s[emb]))));
                }
            }

            return fr;
        });
    }

    getConfig() {
        const config = super.getConfig();
        Object.assign(config, {
            depth: this.clusterSize,
            outputDim: this.n
        });
        return config;
    }
}