View source on GitHub |
An inference sampler that uses a custom sampling function.
Inherits From: Sampler
tfa.seq2seq.InferenceSampler(
sample_fn: Callable,
sample_shape: tfa.types.TensorLike
,
sample_dtype: tfa.types.AcceptableDTypes
,
end_fn: Callable,
next_inputs_fn: Optional[Callable] = None
)
Methods
initialize
initialize(
start_inputs
)
initialize the sampler with the input tensors.
This method must be invoked exactly once before calling other methods of the Sampler.
Args | |
---|---|
inputs
|
A (structure of) input tensors, it could be a nested tuple or a single tensor. |
**kwargs
|
Other kwargs for initialization. It could contain tensors like mask for inputs, or non tensor parameter. |
Returns | |
---|---|
(initial_finished, initial_inputs) .
|
next_inputs
next_inputs(
time, outputs, state, sample_ids
)
Returns (finished, next_inputs, next_state)
.
sample
sample(
time, outputs, state
)
Returns sample_ids
.