tfa.seq2seq.CustomSampler

Base abstract class that allows the user to customize sampling.

Inherits From: Sampler

tfa.seq2seq.CustomSampler(
    initialize_fn: tfa.types.Initializer,
    sample_fn: Callable,
    next_inputs_fn: Callable,
    sample_ids_shape: Optional[TensorLike] = None,
    sample_ids_dtype: tfa.types.AcceptableDTypes = None
)

initialize_fn callable that returns (finished, next_inputs) for the first iteration.
sample_fn callable that takes (time, outputs, state) and emits tensor sample_ids.
next_inputs_fn callable that takes (time, outputs, state, sample_ids) and emits (finished, next_inputs, next_state).
sample_ids_shape Either a list of integers, or a 1-D Tensor of type int32, the shape of each value in the sample_ids batch. Defaults to a scalar.
sample_ids_dtype The dtype of the sample_ids tensor. Defaults to int32.

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 TensorShape. The return value might not available before the invocation of initialize().

Methods

initialize

View source

initialize(
    inputs, **kwargs
)

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

View source

next_inputs(
    time, outputs, state, sample_ids
)

Returns (finished, next_inputs, next_state).

sample

View source

sample(
    time, outputs, state
)

Returns sample_ids.