Python class
RealizationContext
RealizationContext
class max.experimental.tensor.RealizationContext(*args, **kwargs)
Bases: Protocol, AbstractContextManager[RealizationContext]
Implements a way to realize unrealized tensors.
Most users should never have to think about the existence of this type. It exists to facilitate optimizations around where and when tensor operations are executed.
- Each tensor is either real or associated with a RealizationContext.
- If a tensor is not real, ie. “unrealized”, then it is backed by some symbolic computation.
- The RealizationContext is responsible for tracking this symbolic computation and “realizing” the tensor (executing the computation and backing the tensor with real data) if and when it is asked to do so.
- A RealizationContext can only realize tensors associated with it.
RealizationContext abstracts over various semantics of tensor construction.
“Eager” execution: tensors are realized as soon as the realization context exits. This is the default behavior.
This has a huge concrete advantage over eagerly executing one operation at a time: by controlling the boundary of where the eager context starts and ends, we can give advanced users a tool to _enable fine-grained bounds for automatic fusion!
In practice the easiest way to do this is to mark a function as F.functional. This function is then assumed to be “atomic” for the purposes of eager execution. All ops within the function execute as part of the same graph, meaning the compiler is free to fuse operations and generate fused kernels within this region.
“Lazy” execution: tensors are realized only when code later tries to use them.
This enables a class of interface design common in the ML world, in which layers are constructed with randomized weights which are never used. Lazy execution neatly allows constructing entire models, only performing the weight initialization and allocating memory for them if and when those weights are actually used.
Graph compilation: tensors may never be realized.
This allows tensor operations to be composed with direct usage of the Graph API, for instance Module.compile, or using F.* operations in another Graph API usage.
Async execution: Tensors are realized as async functions, allowing clean integration in async systems like web services.
add_source()
add_source(tensor)
Adds a realized tensor as a source of the realization state.
The source is one on whose values unrealized tensors depend.
-
Parameters:
-
tensor (Tensor) – The realized tensor to add as a source to the computation.
-
Returns:
-
A realization state for the tensor. This may be used to compute downstream unrealized values. _If it is used in any mutating operations, it should be assigned to tensor.state to mark the tensor as having been mutated.
-
Return type:
create_unrealized()
create_unrealized(value)
Registers an unrealized graph value with the realization context.
Returns it as an unrealized tensor.
-
Parameters:
-
value (BufferValue | TensorValue) – The graph value representing the result of a computation.
-
Returns:
-
A new tensor associated with the unrealized value.
-
Return type:
graph
graph: Graph
The graph used by the realization context.
realize_all()
async realize_all()
Realizes all unrealized tensors associated with this context.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!