Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets
This commit is contained in:
Kolp
@@ -0,0 +1,529 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import typing
|
||||
from collections.abc import Callable
|
||||
from typing import overload, TYPE_CHECKING, TypeAlias
|
||||
from typing_extensions import ParamSpec, Self, TypeVar
|
||||
|
||||
import torch
|
||||
from torch import Tensor
|
||||
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from torch.xpu import _POOL_HANDLE
|
||||
|
||||
from .._utils import _dummy_type
|
||||
|
||||
|
||||
__all__ = [
|
||||
"is_current_stream_capturing",
|
||||
"graph_pool_handle",
|
||||
"XPUGraph",
|
||||
"graph",
|
||||
"make_graphed_callables",
|
||||
]
|
||||
|
||||
_R = TypeVar("_R")
|
||||
_P = ParamSpec("_P")
|
||||
|
||||
if not hasattr(torch._C, "_XpuStreamBase"):
|
||||
# Define dummy base classes
|
||||
torch._C.__dict__["_XPUGraph"] = _dummy_type("_XPUGraph")
|
||||
torch._C.__dict__["_xpu_graph_pool_handle"] = _dummy_type("_xpu_graph_pool_handle")
|
||||
torch._C.__dict__["_xpu_isCurrentStreamCapturing"] = _dummy_type(
|
||||
"_xpu_isCurrentStreamCapturing"
|
||||
)
|
||||
|
||||
from torch._C import _xpu_graph_pool_handle, _xpu_isCurrentStreamCapturing, _XPUGraph
|
||||
|
||||
|
||||
def is_current_stream_capturing() -> bool:
|
||||
r"""Return True if XPU graph capture is underway on the current XPU stream, False otherwise.
|
||||
|
||||
If a XPU context does not exist on the current device, returns False without initializing the context.
|
||||
"""
|
||||
return _xpu_isCurrentStreamCapturing()
|
||||
|
||||
|
||||
def graph_pool_handle() -> _POOL_HANDLE:
|
||||
r"""Return an opaque token representing the id of a graph memory pool."""
|
||||
return torch.xpu._POOL_HANDLE(_xpu_graph_pool_handle())
|
||||
|
||||
|
||||
class XPUGraph(_XPUGraph):
|
||||
r"""Wrapper around a XPU graph.
|
||||
|
||||
Arguments:
|
||||
keep_graph (bool, optional): If ``keep_graph=False``, the
|
||||
executable command graph will be instantiated on GPU at the end of
|
||||
``capture_end`` and the underlying modifiable command graph will be
|
||||
destroyed. Note that the executable command graph will not be
|
||||
instantiated at the end of ``capture_end`` in this
|
||||
case. Instead, it will be instantiated via an explicit called
|
||||
to ``instantiate`` or automatically on the first call to
|
||||
``replay`` if ``instantiate`` was not already called. Calling
|
||||
``instantiate`` manually before ``replay`` is recommended to
|
||||
prevent increased latency on the first call to ``replay``.
|
||||
|
||||
"""
|
||||
|
||||
def __new__(cls, keep_graph: bool = False) -> Self:
|
||||
return super().__new__(cls, keep_graph)
|
||||
|
||||
def capture_begin(self, pool: _POOL_HANDLE | None = None) -> None:
|
||||
r"""Begin capturing XPU work on the current xpu stream.
|
||||
|
||||
Typically, you shouldn't call ``capture_begin`` yourself.
|
||||
Use :class:`~torch.xpu.graph`, which call ``capture_begin`` internally.
|
||||
|
||||
Arguments:
|
||||
pool (optional): Token (returned by :func:`~torch.xpu.graph_pool_handle` or
|
||||
:meth:`other_Graph_instance.pool()<torch.xpu.XPUGraph.pool>`) that hints this graph may share memory
|
||||
with the indicated pool.
|
||||
"""
|
||||
super().capture_begin(pool=pool)
|
||||
|
||||
def capture_end(self) -> None:
|
||||
r"""End XPU graph capture on the current stream.
|
||||
|
||||
After ``capture_end``, ``replay`` may be called on this instance.
|
||||
|
||||
Typically, you shouldn't call ``capture_end`` yourself.
|
||||
Use :class:`~torch.xpu.graph`, which call ``capture_end`` internally.
|
||||
"""
|
||||
super().capture_end()
|
||||
|
||||
def instantiate(self) -> None:
|
||||
r"""Instantiate the XPU graph. Will be called by
|
||||
``capture_end`` if ``keep_graph=False``, or by ``replay`` if
|
||||
``keep_graph=True`` and ``instantiate`` has not already been
|
||||
explicitly called. Does not destroy the xpu modify command graph returned
|
||||
by ``raw_xpu_graph``.
|
||||
"""
|
||||
super().instantiate()
|
||||
|
||||
def replay(self) -> None:
|
||||
r"""Replay the XPU work captured by this graph."""
|
||||
super().replay()
|
||||
|
||||
def reset(self) -> None:
|
||||
r"""Delete the graph currently held by this instance."""
|
||||
super().reset()
|
||||
|
||||
def pool(self) -> _POOL_HANDLE:
|
||||
r"""Return an opaque token representing the id of this graph's memory pool.
|
||||
|
||||
This id can optionally be passed to another graph's ``capture_begin``,
|
||||
which hints the other graph may share the same memory pool.
|
||||
"""
|
||||
return super().pool()
|
||||
|
||||
def enable_debug_mode(self) -> None:
|
||||
r"""Enable debugging mode for XPUGraph.debug_dump."""
|
||||
return super().enable_debug_mode()
|
||||
|
||||
def debug_dump(self, debug_path: str) -> None:
|
||||
r"""
|
||||
Arguments:
|
||||
debug_path (required): Path to dump the graph to.
|
||||
|
||||
Calls a debugging function to dump the graph if the debugging is
|
||||
enabled via XPUGraph.enable_debug_mode()
|
||||
"""
|
||||
return super().debug_dump(debug_path)
|
||||
|
||||
def raw_xpu_graph(self) -> int:
|
||||
r"""Returns the underlying xpuGraph_t. ``keep_graph`` must be True.
|
||||
|
||||
XPU doesn't provide APIs to manipulate this object.
|
||||
""" # noqa: B950
|
||||
return super().raw_xpu_graph()
|
||||
|
||||
def raw_xpu_graph_exec(self) -> int:
|
||||
r"""Returns the underlying xpuGraphExec_t. ``instantiate`` must have been called if ``keep_graph`` is True, or ``capture_end`` must have been called if ``keep_graph`` is False. If you call ``instantiate()`` after ``raw_xpu_graph_exec()``, the previously returned xpuGraphExec_t will be destroyed. It is your responsibility not to use this object after destruction.
|
||||
|
||||
XPU doesn't provide APIs to manipulate this object.
|
||||
""" # noqa: B950
|
||||
return super().raw_xpu_graph_exec()
|
||||
|
||||
|
||||
class graph:
|
||||
r"""Context-manager that captures XPU work into a :class:`torch.xpu.XPUGraph` object for later replay.
|
||||
|
||||
Arguments:
|
||||
xpu_graph (torch.xpu.XPUGraph): Graph object used for capture.
|
||||
pool (optional): Opaque token (returned by a call to :func:`~torch.xpu.graph_pool_handle()` or
|
||||
:meth:`other_Graph_instance.pool()<torch.xpu.XPUGraph.pool>`) hinting this graph's capture
|
||||
may share memory from the specified pool.
|
||||
stream (torch.xpu.Stream, optional): If supplied, will be set as the current stream in the context.
|
||||
If not supplied, ``graph`` sets its own internal side stream as the current stream in the context.
|
||||
|
||||
.. note::
|
||||
For effective memory sharing, if you pass a ``pool`` used by a previous capture and the previous capture
|
||||
used an explicit ``stream`` argument, you should pass the same ``stream`` argument to this capture.
|
||||
|
||||
""" # noqa: B950
|
||||
|
||||
default_capture_stream: torch.xpu.Stream | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
xpu_graph: XPUGraph,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
stream: torch.xpu.Stream | None = None,
|
||||
):
|
||||
# Lazy-init of default_capture_stream helps avoid circular-import errors.
|
||||
# Not thread safe, but graphs already have the general (explicitly documented)
|
||||
# restriction that only one capture may be underway at a time in the process.
|
||||
if self.__class__.default_capture_stream is None:
|
||||
self.__class__.default_capture_stream = torch.xpu.Stream()
|
||||
|
||||
self.pool: tuple[()] | tuple[_POOL_HANDLE] = () if pool is None else (pool,)
|
||||
self.capture_stream = (
|
||||
stream if stream is not None else self.__class__.default_capture_stream
|
||||
)
|
||||
if self.capture_stream is None:
|
||||
raise AssertionError("capture_stream must not be None")
|
||||
self.stream_ctx = self.capture_stream
|
||||
self.xpu_graph = xpu_graph
|
||||
|
||||
def __enter__(self) -> None:
|
||||
# Free as much memory as we can for the graph
|
||||
torch.xpu.synchronize()
|
||||
|
||||
torch.xpu.empty_cache()
|
||||
self.stream_ctx.__enter__()
|
||||
|
||||
self.xpu_graph.capture_begin(*self.pool)
|
||||
|
||||
def __exit__(self, *args: object) -> None:
|
||||
self.xpu_graph.capture_end()
|
||||
self.stream_ctx.__exit__(*args)
|
||||
|
||||
|
||||
_ModuleOrCallable: TypeAlias = torch.nn.Module | Callable[..., object]
|
||||
|
||||
|
||||
@overload
|
||||
def make_graphed_callables(
|
||||
callables: _ModuleOrCallable,
|
||||
sample_args: tuple[Tensor, ...],
|
||||
num_warmup_iters: int = 3,
|
||||
allow_unused_input: bool = False,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
) -> _ModuleOrCallable: ...
|
||||
|
||||
|
||||
@overload
|
||||
def make_graphed_callables(
|
||||
callables: tuple[_ModuleOrCallable, ...],
|
||||
sample_args: tuple[tuple[Tensor, ...], ...],
|
||||
num_warmup_iters: int = 3,
|
||||
allow_unused_input: bool = False,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
) -> tuple[_ModuleOrCallable, ...]: ...
|
||||
|
||||
|
||||
def make_graphed_callables(
|
||||
callables: _ModuleOrCallable | tuple[_ModuleOrCallable, ...],
|
||||
sample_args: tuple[Tensor, ...] | tuple[tuple[Tensor, ...], ...],
|
||||
num_warmup_iters: int = 3,
|
||||
allow_unused_input: bool = False,
|
||||
pool: _POOL_HANDLE | None = None,
|
||||
) -> _ModuleOrCallable | tuple[_ModuleOrCallable, ...]:
|
||||
r"""Accept callables (functions or :class:`nn.Module<torch.nn.Module>`\ s) and returns graphed versions.
|
||||
|
||||
Each graphed callable's forward pass runs its source callable's
|
||||
forward XPU work as a XPU graph inside a single autograd node.
|
||||
|
||||
The graphed callable's forward pass also appends
|
||||
a backward node to the autograd graph. During backward, this node runs the
|
||||
callable's backward work as a XPU graph.
|
||||
|
||||
Therefore, each graphed callable should be a drop-in replacement for its source callable
|
||||
in an autograd-enabled training loop.
|
||||
|
||||
See :ref:`Partial-network capture<partial-network-capture>` for detailed use and constraints.
|
||||
|
||||
If you pass a tuple of several callables, their captures will use the same memory pool.
|
||||
|
||||
Arguments:
|
||||
callables (torch.nn.Module or Python function, or tuple of these): Callable or callables to graph.
|
||||
If you pass a tuple of callables, their order in the tuple must be the same order they'll run
|
||||
in the live workload.
|
||||
sample_args (tuple of Tensors, or tuple of tuples of Tensors): Samples args for each callable.
|
||||
If a single callable was passed, ``sample_args`` must be a single tuple of argument Tensors.
|
||||
If a tuple of callables was passed, ``sample_args`` must be tuple of tuples of argument Tensors.
|
||||
num_warmup_iters (int): The number of warmup iterations. Currently, ``DataDistributedParallel`` needs
|
||||
11 iterations for warm up. Default: ``3``.
|
||||
allow_unused_input (bool): If False, specifying inputs that were not used when computing outputs
|
||||
(and therefore their grad is always zero) is an error. Defaults to False.
|
||||
pool (optional): Token (returned by :func:`~torch.xpu.graph_pool_handle` or
|
||||
:meth:`other_Graph_instance.pool()<torch.xpu.XPUGraph.pool>`) that hints this graph may share memory
|
||||
with the indicated pool.
|
||||
.. note::
|
||||
The ``requires_grad`` state of each Tensor in ``sample_args`` must match the state
|
||||
that's expected for the corresponding real input in the training loop.
|
||||
|
||||
.. warning::
|
||||
This API is in beta and may change in future releases.
|
||||
|
||||
.. warning::
|
||||
``sample_args`` for each callable must contain only Tensors. Other types are not allowed.
|
||||
|
||||
.. warning::
|
||||
Returned callables do not support higher order differentiation (e.g., double backward).
|
||||
|
||||
.. warning::
|
||||
In any :class:`~torch.nn.Module` passed to :func:`~make_graphed_callables`, only parameters
|
||||
may be trainable. Buffers must have ``requires_grad=False``.
|
||||
|
||||
.. warning::
|
||||
After you pass a :class:`torch.nn.Module` through :func:`~make_graphed_callables`,
|
||||
you may not add or remove any of that Module's parameters or buffers.
|
||||
|
||||
.. warning::
|
||||
:class:`torch.nn.Module`\s passed to :func:`~torch.xpu.make_graphed_callables` must not have module hooks
|
||||
registered on them at the time they are passed. However, registering hooks on modules *after* passing them
|
||||
through :func:`~torch.xpu.make_graphed_callables` is allowed.
|
||||
|
||||
.. warning::
|
||||
When running a graphed callable, you must pass its arguments in the same order and format
|
||||
they appeared in that callable's ``sample_args``.
|
||||
|
||||
.. warning::
|
||||
The automatic mixed precision is supported in :func:`~torch.xpu.make_graphed_callables` only with disabled
|
||||
caching. The context manager `torch.amp.autocast()` must have `cache_enabled=False`.
|
||||
"""
|
||||
if torch.is_autocast_enabled() and torch.is_autocast_cache_enabled():
|
||||
raise RuntimeError(
|
||||
"make_graphed_callables does not support the autocast caching. Please set `cache_enabled=False`."
|
||||
)
|
||||
|
||||
just_one_callable = False
|
||||
|
||||
_sample_args: tuple[tuple[Tensor, ...], ...]
|
||||
if not isinstance(callables, tuple):
|
||||
just_one_callable = True
|
||||
callables = (callables,)
|
||||
_sample_args = (typing.cast(tuple[Tensor, ...], sample_args),)
|
||||
else:
|
||||
_sample_args = typing.cast(tuple[tuple[Tensor, ...], ...], sample_args)
|
||||
|
||||
flatten_sample_args = []
|
||||
|
||||
for c, args in zip(callables, _sample_args):
|
||||
if isinstance(c, torch.nn.Module):
|
||||
if not (
|
||||
len(c._backward_hooks) == 0
|
||||
and len(c._forward_hooks) == 0
|
||||
and len(c._forward_pre_hooks) == 0
|
||||
):
|
||||
raise RuntimeError(
|
||||
"Modules must not have hooks registered at the time they are passed. However, registering hooks "
|
||||
+ "on modules after passing them through make_graphed_callables is allowed."
|
||||
)
|
||||
if not all(b.requires_grad is False for b in c.buffers()):
|
||||
raise RuntimeError(
|
||||
"In any :class:`~torch.nn.Module` passed to "
|
||||
+ ":func:`~make_graphed_callables`, only parameters may be trainable. All buffers must have "
|
||||
+ "``requires_grad=False``."
|
||||
)
|
||||
flatten_arg = torch.utils._pytree.arg_tree_leaves(*args)
|
||||
flatten_sample_args.append(tuple(flatten_arg))
|
||||
if not all(isinstance(arg, torch.Tensor) for arg in flatten_arg):
|
||||
raise TypeError(
|
||||
"In the beta API, sample_args "
|
||||
+ "for each callable must contain only Tensors. Other types are not allowed."
|
||||
)
|
||||
|
||||
# If a callable is an nn.Module, its graph's full input surface is the args the user explicitly
|
||||
# passes to forward (ie, its sample_args) AND the module's parameter attributes.
|
||||
per_callable_len_user_args = [len(args) for args in flatten_sample_args]
|
||||
per_callable_module_params = [
|
||||
tuple(c.parameters()) if isinstance(c, torch.nn.Module) else ()
|
||||
for c in callables
|
||||
]
|
||||
per_callable_static_input_surfaces = [
|
||||
flatten_sample_args[i] + per_callable_module_params[i]
|
||||
for i in range(len(callables))
|
||||
]
|
||||
|
||||
fwd_graphs = [torch.xpu.XPUGraph() for _ in range(len(callables))]
|
||||
bwd_graphs = [torch.xpu.XPUGraph() for _ in range(len(callables))]
|
||||
|
||||
mempool = graph_pool_handle() if pool is None else pool
|
||||
|
||||
# Warmup
|
||||
torch.xpu.synchronize()
|
||||
with torch.xpu.stream(torch.xpu.Stream()):
|
||||
for func, args, static_input_surface in zip(
|
||||
callables, _sample_args, per_callable_static_input_surfaces
|
||||
):
|
||||
grad_inputs, outputs, outputs_grad = None, None, None
|
||||
for _ in range(num_warmup_iters):
|
||||
outputs = torch.utils._pytree.tree_leaves(func(*args))
|
||||
outputs_grad = tuple(o for o in outputs if o.requires_grad)
|
||||
if len(outputs_grad) > 0:
|
||||
grad_inputs = torch.autograd.grad(
|
||||
outputs=outputs_grad,
|
||||
inputs=tuple(
|
||||
i for i in static_input_surface if i.requires_grad
|
||||
),
|
||||
grad_outputs=tuple(
|
||||
torch.empty_like(o) for o in outputs if o.requires_grad
|
||||
),
|
||||
only_inputs=True,
|
||||
allow_unused=allow_unused_input,
|
||||
)
|
||||
for v in [outputs, outputs_grad, grad_inputs]:
|
||||
del v
|
||||
|
||||
torch.xpu.synchronize()
|
||||
|
||||
# Capture forward graphs
|
||||
per_callable_static_outputs = []
|
||||
per_callable_output_unflatten_spec = []
|
||||
for func, args, fwd_graph in zip(callables, _sample_args, fwd_graphs):
|
||||
# each graph uses the same mempool
|
||||
with torch.xpu.graph(fwd_graph, pool=mempool):
|
||||
func_outputs = func(*args)
|
||||
|
||||
flatten_outputs, spec = torch.utils._pytree.tree_flatten(func_outputs)
|
||||
per_callable_static_outputs.append(tuple(flatten_outputs))
|
||||
per_callable_output_unflatten_spec.append(spec)
|
||||
|
||||
# Capture backward graphs in reverse order
|
||||
per_callable_static_grad_outputs = []
|
||||
per_callable_static_grad_inputs = []
|
||||
for static_input_surface, static_outputs, bwd_graph in zip(
|
||||
reversed(per_callable_static_input_surfaces),
|
||||
reversed(per_callable_static_outputs),
|
||||
reversed(bwd_graphs),
|
||||
):
|
||||
static_grad_outputs = tuple(
|
||||
torch.empty_like(o) if o.requires_grad else None for o in static_outputs
|
||||
)
|
||||
|
||||
outputs_grad = tuple(o for o in static_outputs if o.requires_grad)
|
||||
grad_inputs = None
|
||||
if len(outputs_grad) > 0:
|
||||
with torch.xpu.graph(bwd_graph, pool=mempool):
|
||||
grad_inputs = torch.autograd.grad(
|
||||
outputs=outputs_grad,
|
||||
inputs=tuple(i for i in static_input_surface if i.requires_grad),
|
||||
grad_outputs=tuple(o for o in static_grad_outputs if o is not None),
|
||||
only_inputs=True,
|
||||
allow_unused=allow_unused_input,
|
||||
)
|
||||
|
||||
static_grad_inputs = []
|
||||
grad_idx = 0
|
||||
for arg in static_input_surface:
|
||||
if arg.requires_grad and grad_inputs is not None:
|
||||
static_grad_inputs.append(grad_inputs[grad_idx])
|
||||
grad_idx += 1
|
||||
else:
|
||||
static_grad_inputs.append(None) # type: ignore[arg-type]
|
||||
static_grad_inputs = tuple(static_grad_inputs) # type: ignore[assignment]
|
||||
|
||||
per_callable_static_grad_outputs.append(static_grad_outputs)
|
||||
per_callable_static_grad_inputs.append(static_grad_inputs)
|
||||
|
||||
# Reverses the most recent two lists
|
||||
per_callable_static_grad_outputs.reverse()
|
||||
per_callable_static_grad_inputs.reverse()
|
||||
|
||||
def make_graphed_autograd_function(
|
||||
fwd_graph: XPUGraph,
|
||||
bwd_graph: XPUGraph,
|
||||
module_params: tuple[torch.nn.Parameter, ...],
|
||||
len_user_args: int,
|
||||
output_unflatten_spec: torch.utils._pytree.TreeSpec,
|
||||
static_input_surface: tuple[Tensor, ...],
|
||||
static_outputs: tuple[Tensor, ...],
|
||||
static_grad_outputs: tuple[Tensor | None, ...],
|
||||
static_grad_inputs: tuple[Tensor, ...],
|
||||
) -> Callable[..., object]:
|
||||
class Graphed(torch.autograd.Function):
|
||||
@staticmethod
|
||||
# pyrefly: ignore [bad-override]
|
||||
def forward(ctx: object, *inputs: Tensor) -> tuple[Tensor, ...]:
|
||||
# At this stage, only the user args may (potentially) be new tensors.
|
||||
for i in range(len_user_args):
|
||||
if static_input_surface[i].data_ptr() != inputs[i].data_ptr():
|
||||
static_input_surface[i].copy_(inputs[i])
|
||||
fwd_graph.replay()
|
||||
if not isinstance(static_outputs, tuple):
|
||||
raise RuntimeError("static_outputs must be a tuple")
|
||||
return tuple(o.detach() for o in static_outputs)
|
||||
|
||||
@staticmethod
|
||||
@torch.autograd.function.once_differentiable
|
||||
# pyrefly: ignore [bad-override]
|
||||
def backward(ctx: object, *grads: Tensor) -> tuple[Tensor, ...]:
|
||||
if len(grads) != len(static_grad_outputs):
|
||||
raise RuntimeError(
|
||||
f"Expected {len(static_grad_outputs)} gradients but got {len(grads)}"
|
||||
)
|
||||
for g, grad in zip(static_grad_outputs, grads):
|
||||
if g is not None:
|
||||
if g.data_ptr() != grad.data_ptr():
|
||||
g.copy_(grad)
|
||||
bwd_graph.replay()
|
||||
|
||||
if not isinstance(static_grad_inputs, tuple):
|
||||
raise RuntimeError("static_grad_inputs must be a tuple")
|
||||
return tuple(
|
||||
b.detach() if b is not None else b for b in static_grad_inputs
|
||||
)
|
||||
|
||||
def functionalized(*user_args: object) -> object:
|
||||
# Runs the new autograd function which replays the XPU graphs
|
||||
flatten_user_args = torch.utils._pytree.arg_tree_leaves(*user_args)
|
||||
out = Graphed.apply(*(tuple(flatten_user_args) + module_params))
|
||||
return torch.utils._pytree.tree_unflatten(out, output_unflatten_spec)
|
||||
|
||||
return functionalized
|
||||
|
||||
ret: list[_ModuleOrCallable] = []
|
||||
for i, func in enumerate(callables):
|
||||
graphed = make_graphed_autograd_function(
|
||||
fwd_graphs[i],
|
||||
bwd_graphs[i],
|
||||
per_callable_module_params[i],
|
||||
per_callable_len_user_args[i],
|
||||
per_callable_output_unflatten_spec[i],
|
||||
per_callable_static_input_surfaces[i],
|
||||
per_callable_static_outputs[i],
|
||||
per_callable_static_grad_outputs[i],
|
||||
per_callable_static_grad_inputs[i],
|
||||
)
|
||||
|
||||
if isinstance(func, torch.nn.Module):
|
||||
|
||||
def make_graphed_forward(
|
||||
func: torch.nn.Module,
|
||||
graph_training_state: bool,
|
||||
graphed: Callable[_P, _R],
|
||||
orig_fwd: Callable[_P, _R],
|
||||
) -> Callable[_P, _R]:
|
||||
def new_fwd(*user_args: _P.args, **user_kwargs: _P.kwargs) -> _R:
|
||||
if func.training == graph_training_state:
|
||||
return graphed(*user_args, **user_kwargs)
|
||||
else:
|
||||
return orig_fwd(*user_args, **user_kwargs)
|
||||
|
||||
return new_fwd
|
||||
|
||||
func.forward = make_graphed_forward(
|
||||
func, func.training, graphed, func.forward
|
||||
)
|
||||
ret.append(func)
|
||||
else:
|
||||
ret.append(graphed)
|
||||
|
||||
if just_one_callable:
|
||||
return ret[0]
|
||||
|
||||
return tuple(ret)
|
||||
Reference in New Issue
Block a user