(self, input_signature)
Initializes weights and state for inputs with the given signature.
Initializes weights and state for inputs with the given signature.
Initializes weights and state for inputs with the given signature .
def init_weights_and_state(self, input_signature):
    """Initializes weights and state for inputs with the given signature."""
    if self._n_in == 1 and not isinstance(input_signature, (list, tuple)):
      input_signature = (input_signature,)
    batched_signature = []
    for sub_input_signature in input_signature:
      batched_size = 1
      for d in sub_input_signature.shape[:-self._n_last_axes_to_keep]:
        batched_size *= d
      batched_shape = [batched_size] + list(
          sub_input_signature.shape[-self._n_last_axes_to_keep:])
      batched_signature.append(ShapeDtype(batched_shape,
                                          sub_input_signature.dtype))
    if self._n_in == 1:
      batched_signature = batched_signature[0]
    weights, layer_state = self.sublayer.init(batched_signature, use_cache=True)
    self.state = (layer_state,)
    self.weights = (weights,)
def init_weights_and_state ( self , input_signature ) : if self . _n_in == 1 and not isinstance ( input_signature , ( list , tuple ) ) : input_signature = ( input_signature , ) batched_signature = [ ] for sub_input_signature in input_signature : batched_size = 1 for d in sub_input_signature . shape [ : - self . _n_last_axes_to_keep ] : batched_size *= d batched_shape = [ batched_size ] + list ( sub_input_signature . shape [ - self . _n_last_axes_to_keep : ] ) batched_signature . append ( ShapeDtype ( batched_shape , sub_input_signature . dtype ) ) if self . _n_in == 1 : batched_signature = batched_signature [ 0 ] weights , layer_state = self . sublayer . init ( batched_signature , use_cache = True ) self . state = ( layer_state , ) self . weights = ( weights , )