Method test_init

tests/moe_test.py:197–199

(self)

    )

  def test_init(self):
    x = jnp.array([1.0, 2.0]).reshape((1, 1, 2))  # TODO(bug): need reshape due to error
    self.model.init({"params": self.rng, "dropout": self.rng}, x)


class DeepSeekRoutingTest(unittest.TestCase):