neulab/docprompting-conala huggingface dataset:
torch.jit.fork(func, *args, **kwargs) [source]
 
Creates an asynchronous task executing func and a reference to the value of the result of this execution. fork will return immediately, so the return value of func may not have been computed yet. To force completion of the task and access the return value invoke torch.jit.wait on the Future. fork invoked with a func which returns T is typed as torch.jit.Future[T]. fork calls can be arbitrarily nested, and may be invoked with positional and keyword arguments. Asynchronous execution will only occur when run in TorchScript. If run in pure python, fork will not execute in parallel. fork will also not execute in parallel when invoked while tracing, however the fork and wait calls will be captured in the exported IR Graph. .. warning: `fork` tasks will execute non-deterministicly. We recommend only spawning
parallel fork tasks for pure functions that do not modify their inputs,
module attributes, or global state.
  Parameters 
 
func (callable or torch.nn.Module) – A Python function or torch.nn.Module that will be invoked. If executed in TorchScript, it will execute asynchronously, otherwise it will not. Traced invocations of fork will be captured in the IR. 
**kwargs (*args,) – arguments to invoke func with.   Returns 
a reference to the execution of func. The value T can only be accessed by forcing completion of func through torch.jit.wait.  Return type 
torch.jit.Future[T]   Example (fork a free function): import torch
from torch import Tensor
def foo(a : Tensor, b : int) -> Tensor:
    return a + b
def bar(a):
    fut : torch.jit.Future[Tensor] = torch.jit.fork(foo, a, b=2)
    return torch.jit.wait(fut)
script_bar = torch.jit.script(bar)
input = torch.tensor(2)
# only the scripted version executes asynchronously
assert script_bar(input) == bar(input)
# trace is not run asynchronously, but fork is captured in IR
graph = torch.jit.trace(bar, (input,)).graph
assert "fork" in str(graph)
 Example (fork a module method): import torch
from torch import Tensor
class AddMod(torch.nn.Module):
    def forward(self, a: Tensor, b : int):
        return a + b
class Mod(torch.nn.Module):
    def __init__(self):
        super(self).__init__()
        self.mod = AddMod()
    def forward(self, input):
        fut = torch.jit.fork(self.mod, a, b=2)
        return torch.jit.wait(fut)
input = torch.tensor(2)
mod = Mod()
assert mod(input) == torch.jit.script(mod).forward(input)
torch.jit.freeze(mod, preserved_attrs=None, optimize_numerics=True) [source]
 
Freezing a ScriptModule will clone it and attempt to inline the cloned module’s submodules, parameters, and attributes as constants in the TorchScript IR Graph. By default, forward will be preserved, as well as attributes & methods specified in preserved_attrs. Additionally, any attribute that is modified within a preserved method will be preserved. Freezing currently only accepts ScriptModules that are in eval mode.  Parameters 
 
mod (ScriptModule) – a module to be frozen 
preserved_attrs (Optional[List[str]]) – a list of attributes to preserve in addition to the forward method. 
modified in preserved methods will also be preserved. (Attributes) –  
optimize_numerics (bool) – If True, a set of optimization passes will be run that does not strictly 
numerics. Full details of optimization can be found at torch.jit.optimize_frozen_module. (preserve) –    Returns 
Frozen ScriptModule.   Example (Freezing a simple module with a Parameter):     def forward(self, input):
        output = self.weight.mm(input)
        output = self.linear(output)
        return output

scripted_module = torch.jit.script(MyModule(2, 3).eval())
frozen_module = torch.jit.freeze(scripted_module)
# parameters have been removed and inlined into the Graph as constants
assert len(list(frozen_module.named_parameters())) == 0
# See the compiled graph as Python code
print(frozen_module.code)
 Example (Freezing a module with preserved attributes)     def forward(self, input):
        self.modified_tensor += 1
        return input + self.modified_tensor

scripted_module = torch.jit.script(MyModule2().eval())
frozen_module = torch.jit.freeze(scripted_module, preserved_attrs=["version"])
# we've manually preserved `version`, so it still exists on the frozen module and can be modified
assert frozen_module.version == 1
frozen_module.version = 2
# `modified_tensor` is detected as being mutated in the forward, so freezing preserves
# it to retain model semantics
assert frozen_module(torch.tensor(1)) == torch.tensor(12)
# now that we've run it once, the next result will be incremented by one
assert frozen_module(torch.tensor(1)) == torch.tensor(13)
  Note If you’re not sure why an attribute is not being inlined as a constant, you can run dump_alias_db on frozen_module.forward.graph to see if freezing has detected the attribute is being modified.
torch.jit.ignore(drop=False, **kwargs) [source]
 
This decorator indicates to the compiler that a function or method should be ignored and left as a Python function. This allows you to leave code in your model that is not yet TorchScript compatible. If called from TorchScript, ignored functions will dispatch the call to the Python interpreter. Models with ignored functions cannot be exported; use @torch.jit.unused instead. Example (using @torch.jit.ignore on a method): import torch
import torch.nn as nn

class MyModule(nn.Module):
    @torch.jit.ignore
    def debugger(self, x):
        import pdb
        pdb.set_trace()

    def forward(self, x):
        x += 10
        # The compiler would normally try to compile `debugger`,
        # but since it is `@ignore`d, it will be left as a call
        # to Python
        self.debugger(x)
        return x

m = torch.jit.script(MyModule())

# Error! The call `debugger` cannot be saved since it calls into Python
m.save("m.pt")
 Example (using @torch.jit.ignore(drop=True) on a method): import torch
import torch.nn as nn

class MyModule(nn.Module):
    @torch.jit.ignore(drop=True)
    def training_method(self, x):
        import pdb
        pdb.set_trace()

    def forward(self, x):
        if self.training:
            self.training_method(x)
        return x

m = torch.jit.script(MyModule())

# This is OK since `training_method` is not saved, the call is replaced
# with a `raise`.
m.save("m.pt")
torch.jit.isinstance(obj, target_type) [source]
 
This function provides for conatiner type refinement in TorchScript. It can refine parameterized containers of the List, Dict, Tuple, and Optional types. E.g. List[str], Dict[str, List[torch.Tensor]], Optional[Tuple[int,str,int]]. It can also refine basic types such as bools and ints that are available in TorchScript.  Parameters 
 
obj – object to refine the type of 
target_type – type to try to refine obj to   Returns 
 True if obj was successfully refined to the type of target_type,

False otherwise with no new type refinement    Return type 
bool   Example (using torch.jit.isinstance for type refinement): .. testcode: import torch
from typing import Any, Dict, List

class MyModule(torch.nn.Module):
    def __init__(self):
        super(MyModule, self).__init__()

    def forward(self, input: Any): # note the Any type
        if torch.jit.isinstance(input, List[torch.Tensor]):
            for t in input:
                y = t.clamp(0, 0.5)
        elif torch.jit.isinstance(input, Dict[str, str]):
            for val in input.values():
                print(val)

m = torch.jit.script(MyModule())
x = [torch.rand(3,3), torch.rand(4,3)]
m(x)
y = {"key1":"val1","key2":"val2"}
m(y)
torch.jit.is_scripting() [source]
 
Function that returns True when in compilation and False otherwise. This is useful especially with the @unused decorator to leave code in your model that is not yet TorchScript compatible. .. testcode: import torch

@torch.jit.unused
def unsupported_linear_op(x):
    return x

def linear(x):
   if not torch.jit.is_scripting():
      return torch.linear(x)
   else:
      return unsupported_linear_op(x)
torch.jit.load(f, map_location=None, _extra_files=None) [source]
 
Load a ScriptModule or ScriptFunction previously saved with torch.jit.save All previously saved modules, no matter their device, are first loaded onto CPU, and then are moved to the devices they were saved from. If this fails (e.g. because the run time system doesn’t have certain devices), an exception is raised.  Parameters 
 
f – a file-like object (has to implement read, readline, tell, and seek), or a string containing a file name 
map_location (string or torch.device) – A simplified version of map_location in torch.jit.save used to dynamically remap storages to an alternative set of devices. 
_extra_files (dictionary of filename to content) – The extra filenames given in the map would be loaded and their content would be stored in the provided map.   Returns 
A ScriptModule object.   Example: import torch
import io

torch.jit.load('scriptmodule.pt')

# Load ScriptModule from io.BytesIO object
with open('scriptmodule.pt', 'rb') as f:
    buffer = io.BytesIO(f.read())

# Load all tensors to the original device
torch.jit.load(buffer)

# Load all tensors onto CPU, using a device
buffer.seek(0)
torch.jit.load(buffer, map_location=torch.device('cpu'))

# Load all tensors onto CPU, using a string
buffer.seek(0)
torch.jit.load(buffer, map_location='cpu')

# Load with extra files.
extra_files = {'foo.txt': ''}  # values will be replaced with data
torch.jit.load('scriptmodule.pt', _extra_files=extra_files)
print(extra_files['foo.txt'])
torch.jit.save(m, f, _extra_files=None) [source]
 
Save an offline version of this module for use in a separate process. The saved module serializes all of the methods, submodules, parameters, and attributes of this module. It can be loaded into the C++ API using torch::jit::load(filename) or into the Python API with torch.jit.load. To be able to save a module, it must not make any calls to native Python functions. This means that all submodules must be subclasses of ScriptModule as well.  Danger All modules, no matter their device, are always loaded onto the CPU during loading. This is different from torch.load()’s semantics and may change in the future.   Parameters 
 
m – A ScriptModule to save. 
f – A file-like object (has to implement write and flush) or a string containing a file name. 
_extra_files – Map from filename to contents which will be stored as part of f.     Note torch.jit.save attempts to preserve the behavior of some operators across versions. For example, dividing two integer tensors in PyTorch 1.5 performed floor division, and if the module containing that code is saved in PyTorch 1.5 and loaded in PyTorch 1.6 its division behavior will be preserved. The same module saved in PyTorch 1.6 will fail to load in PyTorch 1.5, however, since the behavior of division changed in 1.6, and 1.5 does not know how to replicate the 1.6 behavior.  Example: import torch
import io

class MyModule(torch.nn.Module):
    def forward(self, x):
        return x + 10

m = torch.jit.script(MyModule())

# Save to file
torch.jit.save(m, 'scriptmodule.pt')
# This line is equivalent to the previous
m.save("scriptmodule.pt")

# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.jit.save(m, buffer)

# Save with extra files
extra_files = {'foo.txt': b'bar'}
torch.jit.save(m, 'scriptmodule.pt', _extra_files=extra_files)
class torch.jit.ScriptFunction  
Functionally equivalent to a ScriptModule, but represents a single function and does not have any attributes or Parameters.  
get_debug_state(self: torch._C.ScriptFunction) → torch._C.GraphExecutorState 
  
save(self: torch._C.ScriptFunction, filename: str, _extra_files: Dict[str, str] = {}) → None 
  
save_to_buffer(self: torch._C.ScriptFunction, _extra_files: Dict[str, str] = {}) → bytes
get_debug_state(self: torch._C.ScriptFunction) → torch._C.GraphExecutorState
save(self: torch._C.ScriptFunction, filename: str, _extra_files: Dict[str, str] = {}) → None
save_to_buffer(self: torch._C.ScriptFunction, _extra_files: Dict[str, str] = {}) → bytes
add_module(name, module)  
Adds a child module to the current module. The module can be accessed as an attribute using the given name.  Parameters 
 
name (string) – name of the child module. The child module can be accessed from this module using the given name 
module (Module) – child module to be added to the module.
apply(fn)  
Applies fn recursively to every submodule (as returned by .children()) as well as self. Typical use includes initializing the parameters of a model (see also torch.nn.init).  Parameters 
fn (Module -> None) – function to be applied to each submodule  Returns 
self  Return type 
Module   Example: >>> @torch.no_grad()
>>> def init_weights(m):
>>>     print(m)
>>>     if type(m) == nn.Linear:
>>>         m.weight.fill_(1.0)
>>>         print(m.weight)
>>> net = nn.Sequential(nn.Linear(2, 2), nn.Linear(2, 2))
>>> net.apply(init_weights)
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Linear(in_features=2, out_features=2, bias=True)
Parameter containing:
tensor([[ 1.,  1.],
        [ 1.,  1.]])
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
bfloat16()  
Casts all floating point parameters and buffers to bfloat16 datatype.  Returns 
self  Return type 
Module
buffers(recurse=True)  
Returns an iterator over module buffers.  Parameters 
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.  Yields 
torch.Tensor – module buffer   Example: >>> for buf in model.buffers():
>>>     print(type(buf), buf.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
children()  
Returns an iterator over immediate children modules.  Yields 
Module – a child module
property code  
Returns a pretty-printed representation (as valid Python syntax) of the internal graph for the forward method. See Inspecting Code for details.
property code_with_constants  
Returns a tuple of: [0] a pretty-printed representation (as valid Python syntax) of the internal graph for the forward method. See code. [1] a ConstMap following the CONSTANT.cN format of the output in [0]. The indices in the [0] output are keys to the underlying constant’s values. See Inspecting Code for details.
cpu()  
Moves all model parameters and buffers to the CPU.  Returns 
self  Return type 
Module
cuda(device=None)  
Moves all model parameters and buffers to the GPU. This also makes associated parameters and buffers different objects. So it should be called before constructing optimizer if the module will live on GPU while being optimized.  Parameters 
device (int, optional) – if specified, all parameters will be copied to that device  Returns 
self  Return type 
Module
double()  
Casts all floating point parameters and buffers to double datatype.  Returns 
self  Return type 
Module
eval()  
Sets the module in evaluation mode. This has any effect only on certain modules. See documentations of particular modules for details of their behaviors in training/evaluation mode, if they are affected, e.g. Dropout, BatchNorm, etc. This is equivalent with self.train(False).  Returns 
self  Return type 
Module
extra_repr()  
Set the extra representation of the module To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.
float()  
Casts all floating point parameters and buffers to float datatype.  Returns 
self  Return type 
Module
property graph  
Returns a string representation of the internal graph for the forward method. See Interpreting Graphs for details.
half()  
Casts all floating point parameters and buffers to half datatype.  Returns 
self  Return type 
Module
property inlined_graph  
Returns a string representation of the internal graph for the forward method. This graph will be preprocessed to inline all function and method calls. See Interpreting Graphs for details.
load_state_dict(state_dict, strict=True)  
Copies parameters and buffers from state_dict into this module and its descendants. If strict is True, then the keys of state_dict must exactly match the keys returned by this module’s state_dict() function.  Parameters 
 
state_dict (dict) – a dict containing parameters and persistent buffers. 
strict (bool, optional) – whether to strictly enforce that the keys in state_dict match the keys returned by this module’s state_dict() function. Default: True
   Returns 
 
missing_keys is a list of str containing the missing keys 
unexpected_keys is a list of str containing the unexpected keys   Return type 
NamedTuple with missing_keys and unexpected_keys fields
modules()  
Returns an iterator over all modules in the network.  Yields 
Module – a module in the network    Note Duplicate modules are returned only once. In the following example, l will be returned only once.  Example: >>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.modules()):
        print(idx, '->', m)

0 -> Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
)
1 -> Linear(in_features=2, out_features=2, bias=True)
named_buffers(prefix='', recurse=True)  
Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.  Parameters 
 
prefix (str) – prefix to prepend to all buffer names. 
recurse (bool) – if True, then yields buffers of this module and all submodules. Otherwise, yields only buffers that are direct members of this module.   Yields 
(string, torch.Tensor) – Tuple containing the name and buffer   Example: >>> for name, buf in self.named_buffers():
>>>    if name in ['running_var']:
>>>        print(buf.size())
named_children()  
Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.  Yields 
(string, Module) – Tuple containing a name and child module   Example: >>> for name, module in model.named_children():
>>>     if name in ['conv4', 'conv5']:
>>>         print(module)
named_modules(memo=None, prefix='')  
Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.  Yields 
(string, Module) – Tuple of name and module    Note Duplicate modules are returned only once. In the following example, l will be returned only once.  Example: >>> l = nn.Linear(2, 2)
>>> net = nn.Sequential(l, l)
>>> for idx, m in enumerate(net.named_modules()):
        print(idx, '->', m)

0 -> ('', Sequential(
  (0): Linear(in_features=2, out_features=2, bias=True)
  (1): Linear(in_features=2, out_features=2, bias=True)
))
1 -> ('0', Linear(in_features=2, out_features=2, bias=True))
named_parameters(prefix='', recurse=True)  
Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.  Parameters 
 
prefix (str) – prefix to prepend to all parameter names. 
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.   Yields 
(string, Parameter) – Tuple containing the name and parameter   Example: >>> for name, param in self.named_parameters():
>>>    if name in ['bias']:
>>>        print(param.size())
parameters(recurse=True)  
Returns an iterator over module parameters. This is typically passed to an optimizer.  Parameters 
recurse (bool) – if True, then yields parameters of this module and all submodules. Otherwise, yields only parameters that are direct members of this module.  Yields 
Parameter – module parameter   Example: >>> for param in model.parameters():
>>>     print(type(param), param.size())
<class 'torch.Tensor'> (20L,)
<class 'torch.Tensor'> (20L, 1L, 5L, 5L)
register_backward_hook(hook)  
Registers a backward hook on the module. This function is deprecated in favor of nn.Module.register_full_backward_hook() and the behavior of this function will change in future versions.  Returns 
a handle that can be used to remove the added hook by calling handle.remove()  Return type 
torch.utils.hooks.RemovableHandle
register_buffer(name, tensor, persistent=True)  
Adds a buffer to the module. This is typically used to register a buffer that should not to be considered a model parameter. For example, BatchNorm’s running_mean is not a parameter, but is part of the module’s state. Buffers, by default, are persistent and will be saved alongside parameters. This behavior can be changed by setting persistent to False. The only difference between a persistent buffer and a non-persistent buffer is that the latter will not be a part of this module’s state_dict. Buffers can be accessed as attributes using given names.  Parameters 
 
name (string) – name of the buffer. The buffer can be accessed from this module using the given name 
tensor (Tensor) – buffer to be registered. 
persistent (bool) – whether the buffer is part of this module’s state_dict.    Example: >>> self.register_buffer('running_mean', torch.zeros(num_features))
register_forward_hook(hook)  
Registers a forward hook on the module. The hook will be called every time after forward() has computed an output. It should have the following signature: hook(module, input, output) -> None or modified output
 The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the output. It can modify the input inplace but it will not have effect on forward since this is called after forward() is called.  Returns 
a handle that can be used to remove the added hook by calling handle.remove()  Return type 
torch.utils.hooks.RemovableHandle
register_forward_pre_hook(hook)  
Registers a forward pre-hook on the module. The hook will be called every time before forward() is invoked. It should have the following signature: hook(module, input) -> None or modified input
 The input contains only the positional arguments given to the module. Keyword arguments won’t be passed to the hooks and only to the forward. The hook can modify the input. User can either return a tuple or a single modified value in the hook. We will wrap the value into a tuple if a single value is returned(unless that value is already a tuple).  Returns 
a handle that can be used to remove the added hook by calling handle.remove()  Return type 
torch.utils.hooks.RemovableHandle
register_full_backward_hook(hook)  
Registers a backward hook on the module. The hook will be called every time the gradients with respect to module inputs are computed. The hook should have the following signature: hook(module, grad_input, grad_output) -> tuple(Tensor) or None
 The grad_input and grad_output are tuples that contain the gradients with respect to the inputs and outputs respectively. The hook should not modify its arguments, but it can optionally return a new gradient with respect to the input that will be used in place of grad_input in subsequent computations. grad_input will only correspond to the inputs given as positional arguments and all kwarg arguments are ignored. Entries in grad_input and grad_output will be None for all non-Tensor arguments.  Warning Modifying inputs or outputs inplace is not allowed when using backward hooks and will raise an error.   Returns 
a handle that can be used to remove the added hook by calling handle.remove()  Return type 
torch.utils.hooks.RemovableHandle
register_parameter(name, param)  
Adds a parameter to the module. The parameter can be accessed as an attribute using given name.  Parameters 
 
name (string) – name of the parameter. The parameter can be accessed from this module using the given name 
param (Parameter) – parameter to be added to the module.