Exemplo n.º 1
0
 def initialize(cls, rng, in_spec, initial_state):
   return nn.LayerParams(state=initial_state)
Exemplo n.º 2
0
 def initialize(cls, rng, in_spec):
   return nn.LayerParams()
Exemplo n.º 3
0
 def initialize(cls, rng, in_spec, weight):
   weight = jnp.array(weight)
   return nn.LayerParams(params=weight)