yobx.torch.new_tracing.tensor#

class yobx.torch.new_tracing.tensor.TracingTensor(size: Tuple[int, ...] | TracingShape, dtype: dtype, device: str | device | None = None, requires_grad: bool = False, tracer: GraphTracer | None = None)[source]#

A torch.Tensor subclass that records all dispatch-level operations into a torch.fx.Graph via __torch_dispatch__.

TracingTensor uses __torch_dispatch__ to intercept all tensor operations at the C++ dispatcher level and records them as nodes in a torch.fx.Graph. This produces a full computation graph without requiring Python-level symbolic proxy objects.

Note

TracingTensor instances are created internally by DispatchTracer. Use trace() to trace a callable rather than constructing TracingTensor directly.

It contains two attributes:

  • _tracer: The DispatchTracer managing this tensor’s graph.

  • _node: The torch.fx.Node corresponding to this tensor in the graph.

classmethod from_tensor(t: Tensor, dynamic_shapes: Dict[int, Any] | None = None, tracer: GraphTracer | None = None) → TracingTensor[source]#

Creates a tracing tensor.

make_empty_instance(dyanmic_shape_values: Dict[str, int] | None = None) → Tensor[source]#

Allocates an uninitialised torch.empty() tensor whose dtype and device match this TracingTensor.

Concrete integer dimensions are used as-is. Symbolic (string) dimensions must be resolved by supplying dyanmic_shape_values, a mapping from dimension name to its concrete integer value. A missing entry for any symbolic dimension raises AssertionError.

Parameters:

dyanmic_shape_values – Optional mapping from symbolic dimension names (e.g. "batch") to their concrete integer sizes.

Returns:

A real torch.Tensor with the resolved shape, the same dtype, and the same device as this TracingTensor. The tensor is uninitialised (contents are undefined).

Raises:
property shape: TracingShape#

Returns the shape as a TracingShape.