apache_beam.typehints.pytorch_type_compatibility module¶

class apache_beam.typehints.pytorch_type_compatibility.PytorchBatchConverter(batch_type, element_type, dtype, element_shape=(), partition_dimension=0)[source]¶

Bases: BatchConverter

static from_typehints(element_type, batch_type) → PytorchBatchConverter | None[source]¶

produce_batch(elements)[source]¶

explode_batch(batch)[source]¶: Convert an instance of B to Generator[E].

combine_batches(batches)[source]¶

get_length(batch)[source]¶

estimate_byte_size(batch)[source]¶

class apache_beam.typehints.pytorch_type_compatibility.PytorchTypeHint[source]¶

Bases: object

class PytorchTypeConstraint(dtype, shape=())[source]¶

Bases: TypeConstraint

type_check(batch)[source]¶