def _fast_inference_init_state(input_signature, buffer_length,
                               predict_mask=None):
  """Returns an initial state for causal attention layer fast inference."""
  def zeros_for(batch_size, shape_dtype):
    shape, dtype = shape_dtype.as_tuple()
    d_feature = shape[-1]
    return jnp.zeros((batch_size, buffer_length, d_feature), dtype=dtype)

  batch_size = input_signature[0].shape[0]
  k = zeros_for(batch_size, input_signature[1])
  v = zeros_for(batch_size, input_signature[2])
  if predict_mask is not None:
    mask_for_predict = jnp.zeros((buffer_length,)) != 0
    return (mask_for_predict, k, v, jnp.array(0))
  else:
    return (k, v, jnp.array(0))
Returns an initial state for causal attention layer fast inference.
trax/layers/attention.py