tfp.experimental.auto_batching.type_inference.infer_types_from_signature

Infers the variable types of a given program.

program instructions.Program whose types to infer.
sig A list of (patterns of) instructions.TensorType aligned with program.vars_in.
backend Backend implementation.

typed instructions.Program with types inferred.

ValueError If some types still remain incomplete after inference.