Method get_fake_input

models.py:36–59

(self)

  head_size: int = gin.REQUIRED

  def get_fake_input(self) -> dict[str, Any]:
    fake_input_dict = super().get_fake_input()
    b = self.task_config.batch_size
    n = self.num_heads
    h = self.head_size
    fake_input_dict.update({
        'dstate': tuple(
            [{
                'current_index': jnp.array([0] * b, dtype=jnp.int32),
                'keys': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
                'values': jnp.zeros((b, 2048, n, h), dtype=jnp.bfloat16),
                'recurrent_kvq': None,
                'relative_position_bias': jnp.zeros(
                    (b, n, 1, 1024), dtype=jnp.bfloat16
                ),
            }]
            * 12
        ),
        'eos': jnp.zeros([1024], dtype=jnp.bfloat16),
        'mask': jnp.ones([1024], dtype=jnp.bfloat16),
        'length': 1,
        'temperature': 1.0,
    })
    return fake_input_dict

  def __call__(self, inputs: ...) -> tuple[Any, dict[str, Any]]:
    # Make sure this code is not used on untested cases.