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
)
Args | |
|---|---|
sample_fn
|
A callable that takes outputs and emits tensor
sample_ids.
|
sample_shape
|
Either a list of integers, or a 1-D Tensor of type
int32, the shape of the each sample in the batch returned by
sample_fn.
|
sample_dtype
|
the dtype of the sample returned by sample_fn.
|
end_fn
|
A callable that takes sample_ids and emits a bool vector
shaped [batch_size] indicating whether each sample is an end
token.
|
next_inputs_fn
|
(Optional) A callable that takes sample_ids and
returns the next batch of inputs. If not provided, sample_ids is
used as the next batch of inputs.
|
Attributes | |
|---|---|
batch_size
|
Batch size of tensor returned by sample.
Returns a scalar int32 tensor. The return value might not available before the invocation of initialize(), in this case, ValueError is raised. |
sample_ids_dtype
|
DType of tensor returned by sample.
Returns a DType. The return value might not available before the invocation of initialize(). |
sample_ids_shape
|
Shape of tensor returned by sample, excluding the batch dimension.
Returns a |
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.