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]