yobx.torch.in_torch.patches#

class yobx.torch.in_torch.patches.patched_DynamicDimConstraintPrinter[source]#

Patches torch.tx.experimental.symbolic_shapes.DynamicDimConstraintPrinter._print_Symbol. Valid for torch>=2.10.

yobx.torch.in_torch.patches.patched__broadcast_shapes(*_shapes)[source]#

Patches torch._refs._broadcast_shapes. This patch is needed to export yobx.torch.tiny_models.TinyBroadcastAddModel.

yobx.torch.in_torch.patches.patched__get_range_constraints(mod: Module, export_artifact: ExportArtifact, args, kwargs, dynamic_shapes)[source]#

Patches torch.export._trace._get_range_constraints. See PR #174593.

yobx.torch.in_torch.patches.patched__maybe_broadcast(*args, preserve_cpu_scalar_tensors=True)[source]#

Patches torch._refs._maybe_broadcast. This patch is needed to export yobx.torch.tiny_models.TinyBroadcastAddModel.

yobx.torch.in_torch.patches.patched_infer_size(a, b)[source]#

Patches torch._subclasses.fake_impls.infer_size. This patch is needed to export yobx.torch.tiny_models.TinyBroadcastAddModel.