def test_gumbel_softmax_zero(self):
gumbel_softmax = GumbelSoftmax(temperature=0.01)
logits = randn(4, 6)
gumbels = gumbel_softmax(logits)
print(gumbels)
assert gumbels.shape == (4, 6)
def randn(*shape, dtype=mstype.float32):
return mnp.randn(shape, dtype)
def test_gumbel_softmax_hard(self):
gumbel_softmax = GumbelSoftmax(temperature=0.2, hard=True)
logits = randn(20, 32)
gumbels = gumbel_softmax(logits)
assert gumbels.shape == (20, 32)
def test_bidense(self):
m = BiDense(20, 30, 40)
input1 = mnp.randn(128, 20)
input2 = mnp.randn(128, 30)
output = m(input1, input2)
assert output.shape == (128, 40)