class_docstr: Infinite positional encoding.
class_docstr_tok: ["Infinite","positional","encoding","."]
code: def init_weights_and_state(self, input_signature):
      d_feature = input_signature.shape[-1]
      if self._transform == "diag":
          # Initialize it to a small value because JAX has a bug in softplus.
          scale_isoftplus = (
              jnp.zeros((d_feature,), dtype=jnp.float32) + 1e-4
          )
          weights = scale_isoftplus
      elif self._transform == "any":
          ortho = trax.layers.initializers.OrthogonalInitializer()
          weights = ortho((d_feature, d_feature), self.rng)
      else:
          weights = layer_base.EMPTY_WEIGHTS
      if self._mode == "predict":
          batch_size = input_signature.shape[0]
          self.state = (
              jnp.zeros((batch_size,), dtype=jnp.int32),
              self.rng,
          )
      self.weights = weights
  
imports: import logging
  import jax
  import jax.extend
  import numpy
  import trax
  from trax import fastmath
  from trax.fastmath import numpy
  from trax.layers import base
  from trax.layers import initializers
label_desc: Use image node_id 5 for calling the InfinitePositionalEncoding obj's underlying member method code with example usage: obj.init_weights_and_state(input_signature) without return types