initialization of PHM layers about phc-gnn HOT 4 CLOSED

from benchmarks.utils import set_seed_all
from phc.hypercomplex.layers import PHMLinear
import torch

# Initialize the final weight matrix following a certain distribution
device = "cuda:0" if torch.cuda.is_available() else "cpu"

set_seed_all(seed=43)
phm_lin1 = PHMLinear(in_features=128 // 2, out_features=256 // 2, phm_dim=4, w_init="phm", c_init="standard").to(device)

for w in phm_lin1.W:
    w.data.normal_(mean=0.0, std=0.01)

for w in phm_lin1.W:
    print(w.std())

# tensor(0.0100, device='cuda:0', grad_fn=<StdBackward0>)
# tensor(0.0101, device='cuda:0', grad_fn=<StdBackward0>)
# tensor(0.0099, device='cuda:0', grad_fn=<StdBackward0>)
# tensor(0.0099, device='cuda:0', grad_fn=<StdBackward0>)

for w in phm_lin1.phm_rule:
    w.data.normal_(mean=0.5, std=0.1)

for w in phm_lin1.phm_rule:
    print(w)

# Parameter containing:
# tensor([[0.6034, 0.5514, 0.4601, 0.7307],
#         [0.5802, 0.4613, 0.4960, 0.6374],
#         [0.6922, 0.5066, 0.5063, 0.4360],
#         [0.5713, 0.3694, 0.5513, 0.4803]], device='cuda:0', requires_grad=True)
# Parameter containing:
# tensor([[0.3592, 0.5751, 0.5850, 0.5287],
#         [0.4716, 0.4622, 0.5230, 0.5109],
#         [0.4808, 0.3467, 0.5735, 0.5904],
#         [0.4408, 0.5532, 0.5885, 0.5192]], device='cuda:0', requires_grad=True)
# Parameter containing:
# tensor([[0.3816, 0.6542, 0.3359, 0.4211],
#         [0.6865, 0.3759, 0.5291, 0.5276],
#         [0.6018, 0.5565, 0.4768, 0.6355],
#         [0.5029, 0.5969, 0.6655, 0.3873]], device='cuda:0', requires_grad=True)
# Parameter containing:
# tensor([[0.5919, 0.5583, 0.3676, 0.5180],
#         [0.5897, 0.3686, 0.4941, 0.6941],
#         [0.6832, 0.6234, 0.3679, 0.2792],
#         [0.4790, 0.4572, 0.4511, 0.5616]], device='cuda:0', requires_grad=True)

import torch

from benchmarks.utils import set_seed_all
from phc.hypercomplex.layers import PHMLinear
from phc.hypercomplex.kronecker import kronecker_product_einsum_batched


set_seed_all(42)
phm_dim = 4
in_feats = 256
out_feats = 256
in_feats_axis = in_feats // phm_dim
out_feats_axis = out_feats // phm_dim

# fix this (corresponds to the phm-rules, i.e., the C_i in the paper
C = torch.randn(phm_dim, phm_dim, phm_dim).normal_(0, 0.1)

# try out here
W = torch.randn(phm_dim, in_feats_axis, out_feats_axis).normal_(0, 0.05)

H = kronecker_product_einsum_batched(C, W)
HH = H.sum(0)
print(HH.mean())
print(HH.std())
# tensor(2.9075e-06)
# tensor(0.0087)