https://stats.stackexchange.com/q/592492/69508

592492_69508.py

import tensorflow as tf

class ChiCorr(tf.keras.losses.Loss):

    def chilogl(self, x, k=2):
        ll = k * tf.math.log(x) / 2
        ll = ll - x / 2
        ll = ll - k * tf.math.log(2.0) / 2
        ll = ll - tf.math.lgamma(k / 2)
        return -ll

    def mcorrsq(self, x, y):
        res_x = x - tf.reduce_mean(x)
        res_y = y - tf.reduce_mean(y)
        cov = tf.reduce_mean(res_x * res_y)
        var_x = tf.reduce_mean(res_x**2)
        var_y = tf.reduce_mean(res_y**2)
        sigma_x = tf.sqrt(var_x)
        sigma_y = tf.sqrt(var_y)
        result = cov / (sigma_x * sigma_y)
        result = tf.pow(result, 2.0)
        return -result

    def call(self, y_true, y_pred):
        result = self.chilogl(tf.reduce_sum(y_pred, axis=1), k=2)
        result = result + self.chilogl(y_pred[:,0], k=1)
        result = result + self.mcorrsq(y_pred[:,0], y_pred[:,1])
        return result

class SquaredParameterTensor(tf.keras.models.Model):

    def __init__(self, shape):
        super().__init__()
        self.w = tf.pow(tf.random.normal(shape=shape), 2.0)
        self.w = tf.Variable(self.w, trainable=True)

    def call(self, x=None):
        return self.w

model = SquaredParameterTensor((100000,2))

loss = ChiCorr()
opt = tf.keras.optimizers.Adam(learning_rate=10**-2)

model.compile(loss=loss, optimizer=opt)

history = model.fit(x=[0], y=[0], epochs=10000)

fig, ax = plt.subplots()
ax.scatter(model(None)[:,0], model(None)[:,1])

ax.set_xlabel('X')
ax.set_ylabel('Y')
plt.show()