Python API¶

`freetensor.core` `special` ¶

`freetensor.core.autograd` ¶

`freetensor.core.autograd.grad(func=None, requires=None, provides=None, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, reset_provided_grad=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

Reverse mode automatic differentiation (out-of-place version)

It returns a forward function, a backward function, and two maps on names. The forward function computes the original results. The backward function computes the gradients. The maps map from the names of the original arguments and return values, to the names of their gradients.

If tape_in_closure is True (by default), the forward function has the same interface of the original function, but it will store some intermediate tensors (the tape) in some hidden global states. The backward functions receives the same inputs as the original function plus the graidents of the outputs, and also reads from the hidden states. The outputs of the original function are no longer exist in the backward function, and the input-outputs of the original function are converted to pure inputs. As grad is an out-of-place AD interface, the backward function returns the resulting gradients as additional return values. Names of the additional arguments and return values can be looked up in the maps returned by grad.

If tape_in_closure is False, global states described above will be passed by explicit arguments and return values, so you can store or manipluate these states between the forward run and the backward run.

Parameters:

func (AST) – The original function
requires (Sequence[Union[str, freetensor.core.param_ret_dict.Parameter]]) – Name of input variables that need gradients. A parameter of a function can also be specified with a Parameter object by position
provides (Sequence[Union[str, freetensor.core.param_ret_dict.Parameter, freetensor.core.param_ret_dict.Return]]) – Name of output variables whose gradients are known. A mutable parameter of a function can also be specified with a Parameter object by position. A return value of a function can also be specified with a Return object by position
tapes (Union[Sequence, freetensor_ffi.GradTapeMode, freetensor_ffi.TapeStrategy]) – Intermediate variables that need to be stored from the forward pass and reused in the backward pass. This parameter can be a sequence, which contains VarDef selectors of them. It can also be a GradTapeMode, then it will determine which intermediate variables to be stored by heuristics. Avail GradTapeModes are: All: store all variables including local scalars; None: store nothing; NoReuseOnly: store variables that only hold one version of data, which means we do not have to store each version of them in their history
tape_in_closure (bool) – True to pass taped tensors from the forward function to the backward function in implicit I/O parameters, i.e. in closure. False to pass these tensors as explicit I/O parameters. Default to True
reset_provided_grad (bool) – If true, reset gradients for all variables in provides to 0 after use. This ensures the final result is correct when computing gradients of a program part by part with multiple calls to this function. If false, do not touch the provided gradient, which makes it convenient to run for multiple rounds for timing.
invert (bool) – (Experimental) If set to true, it can reduce the amount of recomputation or taping required. However, this may result in a loss of precision for floating-point numbers.
user_grads (Optional[Sequence[freetensor_ffi.StmtSetToUserGrad]]) – For custom gradient. You do not have to explicitly set this parameter unless you are manipulating func by yourself (not getting it from the Python frontend). See UserGrad for details
attach_backward (bool) – If True, the forward function will be the only return value, with backward function and other metadata attached to it
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
verbose (int) – Verbosity level

Returns:

Func if attach_backward else Tuple – If attach_backward is False, return a tuple.
- Return[0]: forward AST.
- Return[1]: backward AST.
- Return[2]: Mapping from names in requries to its gradient name.
- Return[3]: Mapping from names in provides to its gradient name.
If attach_backward is True, only the forward AST is returned, and others can be get by .backward, .input_name_to_gradient_name and .output_name_to_gradient_name on the return value, respectively.

`freetensor.core.autograd.grad_(func=None, requires=None, provides=None, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, reset_provided_grad=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

Reverse mode automatic differentiation (in-place version)

It returns a forward function, a backward function, and two maps on names. The forward function computes the original results. The backward function computes the gradients. The maps map from the names of the original arguments and return values, to the names of their gradients.

If tape_in_closure is True (by default), the forward function has the same interface of the original function, but it will store some intermediate tensors (the tape) in some hidden global states. The backward functions receives the same inputs as the original function plus the graidents of the outputs, and also reads from the hidden states. The outputs of the original function are no longer exist in the backward function, and the input-outputs of the original function are converted to pure inputs. As grad_ is an in-place AD interface, the backward function passes the resulting gradients by additional mutable arguments. Names of the additional arguments can be looked up in the maps returned by grad_.

If tape_in_closure is False, global states described above will be passed by explicit arguments and return values, so you can store or manipluate these states between the forward run and the backward run.

Parameters:

func (AST) – The original function
requires (Sequence[Union[str, freetensor.core.param_ret_dict.Parameter]]) – Name of input variables that need gradients. A parameter of a function can also be specified with a Parameter object by position
provides (Sequence[Union[str, freetensor.core.param_ret_dict.Parameter, freetensor.core.param_ret_dict.Return]]) – Name of output variables whose gradients are known. A mutable parameter of a function can also be specified with a Parameter object by position. A return value of a function can also be specified with a Return object by position
tapes (Union[Sequence, freetensor_ffi.GradTapeMode, freetensor_ffi.TapeStrategy]) – Intermediate variables that need to be stored from the forward pass and reused in the backward pass. This parameter can be a sequence, which contains VarDef selectors of them. It can also be a GradTapeMode, then it will determine which intermediate variables to be stored by heuristics. Avail GradTapeModes are: All: store all variables including local scalars; None: store nothing; NoReuseOnly: store variables that only hold one version of data, which means we do not have to store each version of them in their history
tape_in_closure (bool) – True to pass taped tensors from the forward function to the backward function in implicit I/O parameters, i.e. in closure. False to pass these tensors as explicit I/O parameters. Default to True
reset_provided_grad (bool) – If true, reset gradients for all variables in provides to 0 after use. This ensures the final result is correct when computing gradients of a program part by part with multiple calls to this function. If false, do not touch the provided gradient, which makes it convenient to run for multiple rounds for timing.
invert (bool) – (Experimental) If set to true, it can reduce the amount of recomputation or taping required. However, this may result in a loss of precision for floating-point numbers.
user_grads (Optional[Sequence[freetensor_ffi.StmtSetToUserGrad]]) – For custom gradient. You do not have to explicitly set this parameter unless you are manipulating func by yourself (not getting it from the Python frontend). See UserGrad for details
attach_backward (bool) – If True, the forward function will be the only return value, with backward function and other metadata attached to it
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
verbose (int) – Verbosity level

Returns:

Func if attach_backward else Tuple – If attach_backward is False, return a tuple.
- Return[0]: forward AST.
- Return[1]: backward AST.
- Return[2]: Mapping from names in requries to its gradient name.
- Return[3]: Mapping from names in provides to its gradient name.
If attach_backward is True, only the forward AST is returned, and others can be get by .backward, .input_name_to_gradient_name and .output_name_to_gradient_name on the return value, respectively.

`freetensor.core.autograd.grad_body(stmt, requires, provides, tapes=<GradTapeMode.NoReuseOnly: 2>, reset_provided_grad=True, invert=False, user_grads=[])` ¶

grad or grad_ on a function body (for internal tests only)

`freetensor.core.autograd.jacrev(func=None, inputs=None, output=None, flatten=False, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

Compute Jacobian tensors using Reverse mode automatic differentiation (out-of-place)

jacrev computes one Jacobian tensor for one output and one or more inputs. Each Jacobian tensor consists of derivatives of all elements in the output tensor w.r.t. all elements in each inputs tensor.

It returns a forward function, a backward function, and a map on names. The forward function computes the original results. The backward function computes the Jacobian tensors. The map maps from the names of the original arguments to the names of their Jacobian tensors.

By default, the forward function has the same interface of the original function, but it will store some intermediate tensors (the tape) in some hidden global states. The backward functions receives the same inputs as the original function, and also reads from the hidden states. The outputs of the original function are no longer exist in the backward function, and the input-outputs of the original function are converted to pure inputs. As jacrev is an out-of-place interface, the backward function returns the resulting Jacobian as additional return values. Names of the additional return values can be looked up in the map returned by jacrev.

Suppose the output's shape is (d1, d2, ...), and there are two inputs, whose shapes are (e1, e2, ...) and (f1, f2, ...), respectively. If flatten is False (by default), the Jacobian tensors' shape will be (d1, d2, ..., e1, e2, ...) and (d1, d2, ..., f1, f2, ...), respectively. If flatten is True, there will be only one Jacbian tensor, whose shape will be (d1 * d2 * ..., e1 * e2 * ... + f1 * f2 * ...).

If tape_in_closure is False, global states described above will be passed by explicit arguments and return values, so you can store or manipluate these states between the forward run and the backward run.

Parameters:

func (AST) – The original function
inputs (Sequence[Union[str, freetensor.core.param_ret_dict.Parameter]]) – Name of input variables that the Jacobian tensors are for.
output (Union[str, freetensor.core.param_ret_dict.Parameter, freetensor.core.param_ret_dict.Return]) – Name of one output variables that the Jacobian tensors are for. A return value of a function can be specified with a Return object
flatten (bool) – If True, concatenate all Jacobian tensors together, to form an (n, m)-shaped output, where n is the total number of elements in the specified output, and m is the total number of elements in the specified inputs. This requires all involved inputs having the same data type and memory type. In this case, the name of the Jacobian tensor will be "jacrev.flatten", and the returned name map will be empty
tapes (Union[Sequence, freetensor_ffi.GradTapeMode, freetensor_ffi.TapeStrategy]) – Intermediate variables that need to be stored from the forward pass and reused in the backward pass. This parameter can be a sequence, which contains VarDef selectors of them. It can also be a GradTapeMode, then it will determine which intermediate variables to be stored by heuristics. Avail GradTapeModes are: All: store all variables including local scalars; None: store nothing; NoReuseOnly: store variables that only hold one version of data, which means we do not have to store each version of them in their history
tape_in_closure (bool) – True to pass taped tensors from the forward function to the backward function in implicit I/O parameters, i.e. in closure. False to pass these tensors as explicit I/O parameters. Default to True
invert (bool) – (Experimental) If set to true, it can reduce the amount of recomputation or taping required. However, this may result in a loss of precision for floating-point numbers.
user_grads (Optional[Sequence[freetensor_ffi.StmtSetToUserGrad]]) – For custom gradient. You do not have to explicitly set this parameter unless you are manipulating func by yourself (not getting it from the Python frontend). See UserGrad for details
attach_backward (bool) – If True, the forward function will be the only return value, with backward function and other metadata attached to it
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
verbose (int) – Verbosity level

Returns:	`Func if attach_backward else Tuple` – If `attach_backward` is False, return a tuple. Return[0]: forward AST. Return[1]: backward AST. Return[2]: Mapping from names in inputs to its Jacobian tensor name. If `attach_backward` is True, only the forward AST is returned, and others can be get by `.backward` and `.input_name_to_gradient_name` on the return value, respectively.

`freetensor.core.autograd.jacrev_(func=None, inputs=None, output=None, flatten=False, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

Compute Jacobian tensors using Reverse mode automatic differentiation (in-place)

jacrev computes one Jacobian tensor for one output and one or more inputs. Each Jacobian tensor consists of derivatives of all elements in the output tensor w.r.t. all elements in each inputs tensor.

It returns a forward function, a backward function, and a map on names. The forward function computes the original results. The backward function computes the Jacobian tensors. The map maps from the names of the original arguments to the names of their Jacobian tensors.

By default, the forward function has the same interface of the original function, but it will store some intermediate tensors (the tape) in some hidden global states. The backward functions receives the same inputs as the original function, and also reads from the hidden states. The outputs of the original function are no longer exist in the backward function, and the input-outputs of the original function are converted to pure inputs. As jacrev_ is an in-place interface, the backward function passes the resulting gradients by additional mutable arguments. Names of the additional arguments can be looked up in the map returned by jacrev_.

Suppose the output's shape is (d1, d2, ...), and there are two inputs, whose shapes are (e1, e2, ...) and (f1, f2, ...), respectively. If flatten is False (by default), the Jacobian tensors' shape will be (d1, d2, ..., e1, e2, ...) and (d1, d2, ..., f1, f2, ...), respectively. If flatten is True, there will be only one Jacbian tensor, whose shape will be (d1 * d2 * ..., e1 * e2 * ... + f1 * f2 * ...).

If tape_in_closure is False, global states described above will be passed by explicit arguments and return values, so you can store or manipluate these states between the forward run and the backward run.

Parameters:

func (AST) – The original function
inputs (Sequence[Union[str, freetensor.core.param_ret_dict.Parameter]]) – Name of input variables that the Jacobian tensors are for. A parameter of a function can also be specified with a Parameter object by position
output (Union[str, freetensor.core.param_ret_dict.Parameter, freetensor.core.param_ret_dict.Return]) – Name of one output variables that the Jacobian tensors are for. A mutable parameter of a function can also be specified with a Parameter object by position. A return value of a function can also be specified with a Return object by position
flatten (bool) – If True, concatenate all Jacobian tensors together, to form an (n, m)-shaped output, where n is the total number of elements in the specified output, and m is the total number of elements in the specified inputs. This requires all involved inputs having the same data type and memory type. In this case, the name of the Jacobian tensor will be "jacrev.flatten", and the returned name map will be empty
tapes (Union[Sequence, freetensor_ffi.GradTapeMode]) – Intermediate variables that need to be stored from the forward pass and reused in the backward pass. This parameter can be a sequence, which contains VarDef selectors of them. It can also be a GradTapeMode, then it will determine which intermediate variables to be stored by heuristics. Avail GradTapeModes are: All: store all variables including local scalars; None: store nothing; NoReuseOnly: store variables that only hold one version of data, which means we do not have to store each version of them in their history
tape_in_closure (bool) – True to pass taped tensors from the forward function to the backward function in implicit I/O parameters, i.e. in closure. False to pass these tensors as explicit I/O parameters. Default to True
invert (bool) – (Experimental) If set to true, it can reduce the amount of recomputation or taping required. However, this may result in a loss of precision for floating-point numbers.
user_grads (Optional[Sequence[freetensor_ffi.StmtSetToUserGrad]]) – For custom gradient. You do not have to explicitly set this parameter unless you are manipulating func by yourself (not getting it from the Python frontend). See UserGrad for details
attach_backward (bool) – If True, the forward function will be the only return value, with backward function and other metadata attached to it
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
verbose (int) – Verbosity level

Returns:	`Func if attach_backward else Tuple` – If `attach_backward` is False, return a tuple. Return[0]: forward AST. Return[1]: backward AST. Return[2]: Mapping from names in inputs to its Jacobian tensor name. If `attach_backward` is True, only the forward AST is returned, and others can be get by `.backward` and `.input_name_to_gradient_name` on the return value, respectively.

`freetensor.core.codegen` ¶

`freetensor.core.codegen.NativeCode (EnableAttachBackward, NativeCode)` ¶

Generated native code with metadata

NOTE: This class does not support serialization yet. If you need serialization, serialize the Func, and re-run codegen.

`freetensor.core.codegen.NativeCode.contains(self, item)` `special` ¶

Legacy interface for testing if a string is in the code

`freetensor.core.codegen.codegen(ast=None, target=None, jit_cache=<function cache at 0x147b6469ca60>, verbose=False)` ¶

Generate native code

Parameters:

ast (Func) – The AST to be lowered. It must includes function signature to determine parameters and return values. If not specified, a partial function is returned, which can be used as a decorator
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
target (Optional[freetensor_ffi.Target]) – The target architecture. If omitted, use the default one in config

Returns:	`NativeCode` – Return a NativeCode for the generated code if there is no JIT parameters. Return a JITTemplate that generates a NativeCode if there is at least one

`freetensor.core.config` ¶

Global configurations

`freetensor.core.config.backend_compiler_cxx(*args, **kvs)` ¶

backend_compiler_cxx() -> List[str]

Backend compiler used to compile generated C++ code

`freetensor.core.config.backend_compiler_nvcc(*args, **kvs)` ¶

backend_compiler_nvcc() -> List[str]

Backend compiler used to compile generated CUDA code

`freetensor.core.config.backend_openmp(*args, **kvs)` ¶

backend_openmp() -> List[str]

OpenMP library linked to the compiled executable

`freetensor.core.config.debug_binary(*args, **kvs)` ¶

debug_binary() -> bool

Check if compiling binary in debug mode

`freetensor.core.config.debug_cuda_with_um(*args, **kvs)` ¶

debug_cuda_with_um() -> bool

Check if debugging with Unified Memory enabled

`freetensor.core.config.default_device(*args, **kvs)` ¶

default_device() -> freetensor_ffi.Device

Check current default device

`freetensor.core.config.default_target(*args, **kvs)` ¶

default_target() -> freetensor_ffi.Target

Check current default target

`freetensor.core.config.fast_math(*args, **kvs)` ¶

fast_math() -> bool

Run pass/float_simplify optimization pass, and enable fast math on backend compilers

`freetensor.core.config.pretty_print(*args, **kvs)` ¶

pretty_print() -> bool

Check if colored printing enabled

`freetensor.core.config.print_all_id(*args, **kvs)` ¶

print_all_id() -> bool

Check if printing IDs of all statements in an AST

`freetensor.core.config.print_source_location(*args, **kvs)` ¶

print_source_location() -> bool

Check if printing Python source location of all statements in an AST

`freetensor.core.config.runtime_dir(*args, **kvs)` ¶

runtime_dir() -> List[str]

Serach paths for FreeTensor runtime header files

`freetensor.core.config.set_backend_compiler_cxx(*args, **kvs)` ¶

set_backend_compiler_cxx(path: List[str]) -> None

Set backend compiler used to compile generated C++ code, unescaped raw path expected

`freetensor.core.config.set_backend_compiler_nvcc(*args, **kvs)` ¶

set_backend_compiler_nvcc(path: List[str]) -> None

Set backend compiler used to compile generated CUDA code, unescaped raw path expected

`freetensor.core.config.set_backend_openmp(*args, **kvs)` ¶

set_backend_openmp(arg0: List[str]) -> None

Set the OpenMP library linked to the compiled executable

`freetensor.core.config.set_debug_binary(*args, **kvs)` ¶

set_debug_binary(flag: bool = True) -> None

Compile with -g at backend. FreeTensor will not delete the binary file after loading it

`freetensor.core.config.set_debug_cuda_with_um(*args, **kvs)` ¶

set_debug_cuda_with_um(arg0: bool) -> None

Allocate CUDA buffers on Unified Memory, for faster (debugging) access of GPU Array from CPU, but with slower Array allocations and more synchronizations. No performance effect on normal in-kernel computations

`freetensor.core.config.set_default_device(*args, **kvs)` ¶

set_default_device(device: freetensor_ffi.Device) -> None

Set default device (internal implementation of with Device)

`freetensor.core.config.set_default_target(*args, **kvs)` ¶

set_default_target(target: freetensor_ffi.Target) -> None

Set default target (internal implementation of with Target)

`freetensor.core.config.set_fast_math(*args, **kvs)` ¶

set_fast_math(flag: bool = True) -> None

Set to run pass/float_simplify optimization pass, and enable fast math on backend compilers (or not)

`freetensor.core.config.set_pretty_print(*args, **kvs)` ¶

set_pretty_print(flag: bool = True) -> None

Set colored printing

`freetensor.core.config.set_print_all_id(*args, **kvs)` ¶

set_print_all_id(flag: bool = True) -> None

Print IDs of all statements in an AST

`freetensor.core.config.set_print_source_location(*args, **kvs)` ¶

set_print_source_location(flag: bool = True) -> None

Print Python source location of all statements in an AST

`freetensor.core.config.set_runtime_dir(*args, **kvs)` ¶

set_runtime_dir(paths: List[str]) -> None

Set serach paths for FreeTensor runtime header files

`freetensor.core.config.set_werror(*args, **kvs)` ¶

set_werror(flag: bool = True) -> None

Error on warning

`freetensor.core.config.werror(*args, **kvs)` ¶

werror() -> bool

Check if error-on-warning enabled

`freetensor.core.config.with_cuda(*args, **kvs)` ¶

with_cuda() -> bool

Check if FreeTensor is built with CUDA

`freetensor.core.config.with_mkl(*args, **kvs)` ¶

with_mkl() -> bool

Check if FreeTensor is built with MKL

`freetensor.core.config.with_pytorch(*args, **kvs)` ¶

with_pytorch() -> bool

Check if FreeTensor is built with PyTorch interface

`freetensor.core.context` ¶

Facility to pick statements to build an AST

Classes and functions in this module are internally used by transformer to construct ASTs. They are also used by some internal tests. API of these classes and functions are subject to changes. End users are encouraged to use transformer, instead of this module.

`freetensor.core.context.ContextStack` ¶

`freetensor.core.context.ContextStack.get_last_stmt_id(self)` ¶

Can be used inside the staged code, to get the ID of the immediately preceding statement

`freetensor.core.context.ContextStack.push_append_stmt_callback(self, callback)` ¶

Add a callback to be called with all next statements to be appended. For If statement, it

can be called twice, one without "else" branch, and then maybe one more with "else" branch

`freetensor.core.context.StmtRange` ¶

Record a set of statement in a program, can be used for custom gradient

Usage:

with StmtRange() as rng:
    # Some statements

StmtRange can be used interleaved with AST scopes. In these cases, you can directly call __enter__ and __exit__. E.g.,

rng = StmtRange()
rng.__enter__()
# Some statements
with VarDef(...)  # Some scopes
    # Some other statements
    rng.__exit__(None, None, None)

`freetensor.core.context.pop_ast(verbose=False)` ¶

Get AST and reset context

Internally used by transformer and tests

`freetensor.core.context.pop_ast_and_user_grads(verbose=False)` ¶

Get AST and reset context. Return an extra list for custom gradients

Set UserGrad for details

`freetensor.core.driver` ¶

`freetensor.core.driver.Driver (EnableAttachBackward, Driver)` ¶

`freetensor.core.driver.Driver.call(self, *args, **kws)` `special` ¶

Set argument, execute the binary code, and collect the returns

If there is only one return value, it is returned directly. Otherwise, the return values are packed in a ReturnValuesPack

This function will introduce some overhaed handling arguments and return values. For an accurate execution time measurement, plase call self.set_args first, then self.time, and finally self.collect_returns

`freetensor.core.driver.Driver.init(self, native_code, device=None, host_device=None, cxx_flags=[], verbose=False)` `special` ¶

Compile a program using a backend compiler and load it into memory

This class is for internal use. Please consider using build_binary

Parameters:	native_code (`NativeCode`) – Native code generated from codegen device (`Optional[freetensor_ffi.Device]`) – The device to run the program. If omitted, use the default device in config cxx_flags (`Sequence[str]`) – Additional C++ flags passed to the backend compiler verbose (`bool`) – True to print extra infomation

`freetensor.core.driver.Driver.collect_returns(self, always_return_pack=False)` ¶

Collect return values from an invocation

Return values must be collect. Otherwise there will be memory leaks

If there is only one return value, it is returned directly. Otherwise, or if always_return_pack is set, the return values are packed in a ReturnValuesPack

`freetensor.core.driver.Driver.native_code(self)` ¶

Get native code compiled by backend compiler

`freetensor.core.driver.Driver.set_args(self, *args, **kws)` ¶

Set argument for an invocation

`freetensor.core.driver.array(data, dtype=None, dont_drop_borrow=False, moved=False)` ¶

Factory function for Array

This function is preferred over directly calling Array's constructor, because it accepts more data format.

If data is another FreeTensor Array, the original object will be returned, with dont_drop_borrow and moved set to new values. If dtype is set and different from the original data type, the Array will be copied first to convert the data type.
If data is Numpy Array or PyTorch Tensor, it will be converted to FreeTensor Array. Memory copy will be avoided in most cases, but it is inevitable if the data is strided. If dtype is set and different from the original data type, the Array or Tensor will be copied first to convert the data type.
Otherwise, the data will be treated as an n-dimensional array-like object, and will be parsed according the rules in NumPy. The data type is also set accordingly, unless dtype is set.

Parameters:

data (FreeTensor Array, Numpy Array, PyTorch Tensor, or other array-like objects) – Data to be copied to or borrowed by the new Array object
dtype (ft.DataType or str) – If data is not in dtype, convert it to dtype first before constructing the Array
dont_drop_borrow (bool) – If true, report an error if we have to drop a borrwed data. This flag is set to true when the Array is cunstructed IMPLICITLY (not by this function) from a user object by borrowing from it, where users may expect they are acutually manipulating the their user object, instead of this Array
moved (bool) – If true, it means we do not care about data in this Array any more after the program runs. Variables with "input-mutable" access type may modify the Array

`freetensor.core.driver.build_binary(code=None, device=None, host_device=None, jit_cache=<function cache at 0x147b6469ca60>, cxx_flags=[], verbose=False)` ¶

Compile a program using a backend compiler and load it into memory

Parameters:

code (Optional[freetensor.core.codegen.NativeCode]) – Native code generated by codegen. If not specified, a partial function is returned, which can be used as a decorator
device (Optional[freetensor_ffi.Device]) – The device to run the program. If omitted, use the default device in config
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
cxx_flags (Sequence[str]) – Additional C++ flags passed to the backend compiler
verbose (bool) – Verbosity level

Returns:	`Driver or JITTemplate` – Return a Driver for the executable if there is no JIT parameters. Return a JITTemplate that generates a Driver if there is at least one

`freetensor.core.driver.move(data)` ¶

Alias for array(data, dont_drop_borrow=False, moved=True)

`freetensor.core.enable_attach_backward` ¶

`freetensor.core.enable_attach_backward.EnableAttachBackward` ¶

Get backward object (Func, Driver, etc) and other meta data from a forward object

This class is a Mixin Class. It should be inherited BEFORE other base classes in multiple inheritance.

`freetensor.core.enable_attach_backward.EnableAttachBackward.init(self, *args, **kvs)` `special` ¶

Forward all arguments to other base classes

In Python, super().init calls the next base class in the full inheritance graph of the final class, not only base classes of BackwardAttachedMixin. See https://docs.python.org/3/tutorial/classes.html#multiple-inheritance

`freetensor.core.expr` ¶

Facility to build AST expressions

Classes and functions in this module are not only used internally for constructing AST nodes, and also exposed to users via multi-stage programming

`freetensor.core.expr.AlreadyMadeReduceTo` ¶

A single-value type that marks a ReduceTo node is already made, and there is no need to

make another Store node

In standard Python data model, functions like iadd returns the modified self, and setitem does a self-assignment. We do the augmenting assignment directly in iadd and return AlreadyMadeReduceTo, so we do not have to Store it again

`freetensor.core.expr.UndeclaredParam` `dataclass` ¶

Error type. For error reporting only.

`freetensor.core.expr.VarRef (FrontendVar)` ¶

Variable of FreeTensor

All variables in FreeTensor DSL (declared via Var, created by empty or var, returned by libop, etc.), and their slices, are VarRef objects. Operations on VarRef objects generates AST nodes

`freetensor.core.expr.VarRef.as_reduce_to(self, reduce_op, metadata, value, atomic=False)` ¶

as_reduce_to(self: freetensor_ffi.FrontendVar, op: freetensor_ffi.ReduceOp, metadata: freetensor_ffi.Metadata, value: freetensor_ffi.Expr, atomic: bool = False) -> freetensor_ffi.Stmt

`freetensor.core.expr.VarRef.as_store(self, metadata, value)` ¶

as_store(self: freetensor_ffi.FrontendVar, metadata: freetensor_ffi.Metadata, value: freetensor_ffi.Expr) -> freetensor_ffi.Stmt

`freetensor.core.expr.VarRef.select(self, idx, dim)` ¶

Alias for self[..., idx, ...], where idx is at the dim-th dimension

`freetensor.core.expr.VarRef.select_slice(self, begin, end, *, dim)` ¶

Alias for self[..., begin:end, ...], where begin:end is at the dim-th dimension

`freetensor.core.expr.VarRef.shape(self, dim=None)` ¶

Return lengths of all dimensions or the length of one dimension

.shape() -> list of lengths of all dimensions

.shape(dim) -> length of dimension dim, where dim can be int or Expr

All lengths can be Expr (if the length is dynamically decided) or int (if statically decided)

`freetensor.core.expr.VarRefFromVarDef (VarRef)` ¶

VarRef with extra checks

`freetensor.core.expr.VarRefFromVarDef.as_reduce_to(self, reduce_op, metadata, value, atomic=False)` `inherited` ¶

as_reduce_to(self: freetensor_ffi.FrontendVar, op: freetensor_ffi.ReduceOp, metadata: freetensor_ffi.Metadata, value: freetensor_ffi.Expr, atomic: bool = False) -> freetensor_ffi.Stmt

`freetensor.core.expr.VarRefFromVarDef.as_store(self, metadata, value)` `inherited` ¶

as_store(self: freetensor_ffi.FrontendVar, metadata: freetensor_ffi.Metadata, value: freetensor_ffi.Expr) -> freetensor_ffi.Stmt

`freetensor.core.expr.VarRefFromVarDef.select(self, idx, dim)` `inherited` ¶

Alias for self[..., idx, ...], where idx is at the dim-th dimension

`freetensor.core.expr.VarRefFromVarDef.select_slice(self, begin, end, *, dim)` `inherited` ¶

Alias for self[..., begin:end, ...], where begin:end is at the dim-th dimension

`freetensor.core.expr.VarRefFromVarDef.shape(self, dim=None)` `inherited` ¶

Return lengths of all dimensions or the length of one dimension

.shape() -> list of lengths of all dimensions

.shape(dim) -> length of dimension dim, where dim can be int or Expr

All lengths can be Expr (if the length is dynamically decided) or int (if statically decided)

`freetensor.core.expr.VarVersionRef (VarRef)` ¶

Special VarRef used for custom gradient, generated from mark_version

`freetensor.core.expr.VarVersionRef.as_reduce_to(self, reduce_op, metadata, value, atomic=False)` `inherited` ¶

as_reduce_to(self: freetensor_ffi.FrontendVar, op: freetensor_ffi.ReduceOp, metadata: freetensor_ffi.Metadata, value: freetensor_ffi.Expr, atomic: bool = False) -> freetensor_ffi.Stmt

`freetensor.core.expr.VarVersionRef.as_store(self, metadata, value)` `inherited` ¶

as_store(self: freetensor_ffi.FrontendVar, metadata: freetensor_ffi.Metadata, value: freetensor_ffi.Expr) -> freetensor_ffi.Stmt

`freetensor.core.expr.VarVersionRef.select(self, idx, dim)` `inherited` ¶

Alias for self[..., idx, ...], where idx is at the dim-th dimension

`freetensor.core.expr.VarVersionRef.select_slice(self, begin, end, *, dim)` `inherited` ¶

Alias for self[..., begin:end, ...], where begin:end is at the dim-th dimension

`freetensor.core.expr.VarVersionRef.shape(self, dim=None)` `inherited` ¶

Return lengths of all dimensions or the length of one dimension

.shape() -> list of lengths of all dimensions

.shape(dim) -> length of dimension dim, where dim can be int or Expr

All lengths can be Expr (if the length is dynamically decided) or int (if statically decided)

`freetensor.core.expr.abs(expr)` ¶

Absolute value

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.abs

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The absolute value

`freetensor.core.expr.add(lhs, rhs)` ¶

lhs + rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.add

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The sum

`freetensor.core.expr.any()` ¶

Create an AnyExpr node (only for testing and type inference)

Any nodes matches any expression nodes in ast.match

`freetensor.core.expr.cast(expr, dtype)` ¶

Cast to another type

Parameters:	expr (`VarRef or Number`) – The operand dtype (`DataTypr or str`) – The target data type

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.ceil(expr)` ¶

Round a float up to an interger (towards +inf)

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.ceil

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.ceildiv(lhs, rhs)` ¶

Ceiling integer division of lhs dividing by rhs

The result rounds towards positive infinity

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.ceildiv

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The quotient

`freetensor.core.expr.cos(expr)` ¶

Cosine

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.tanh

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.dtype(var)` ¶

Get element data type of a variable

`freetensor.core.expr.eq(lhs, rhs)` ¶

lhs == rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.eq

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The comparison

`freetensor.core.expr.exp(expr)` ¶

Natural exponent

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.exp

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The exponent

`freetensor.core.expr.floor(expr)` ¶

Round a float down to an interger (towards -inf)

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.floor

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.floordiv(lhs, rhs)` ¶

Floored integer division of lhs dividing by rhs

The result rounds towards negative infinity (following Python convention, instead of C) This function is recommended over round_towards_0_div, as it enjoys more optimizations

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.floordiv

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The quotient

`freetensor.core.expr.ge(lhs, rhs)` ¶

lhs >= rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.ge

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The comparison

`freetensor.core.expr.gt(lhs, rhs)` ¶

lhs > rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.gt

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The comparison

`freetensor.core.expr.if_then_else(cond, then_case, else_case)` ¶

Similar to then_case if cond else else_case

NOTE: there is NO guarantee that only one branch will be executed. In some cases, both branches will be executed and the result of one of them will be picked. Therefore, please do NOT use if_then_else to guard an out-of-bound array indexing

Parameters:	cond (`VarRef of Number`) – Condition then_case (`VarRef or Number`) – Then-case experssion else_case (`VarRef or Number`) – Else-case expression

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.intrinsic(fmt, params, , ret_type='void', has_side_effect=False)` ¶

Invoke whatever target code

Parameters:

fmt (str) – What to run. "%" is filled by parameters one by one. E.g. sinf(%). Use "%%" to escape for "%". If you need two adjacent parameters, type "(%)(%)" or "% %".
*params (Sequence[Expr]) – (Positional variadic) Parameters to fmt
ret_type (DataType or str) – (Keyword argument only) The return type. Void for no return type. Defaults to Void
has_side_effect (bool) – (Keyword argument only) True to indicate the intrinsic modifes something other than the return value. Defaults to false

`freetensor.core.expr.l_and(lhs, rhs)` ¶

Logical and of lhs and rhs

NOTE: Short-circuit evaluation is NOT supported

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.l_and

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The logical and

`freetensor.core.expr.l_not(expr)` ¶

Logical not

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.l_not

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The logical not

`freetensor.core.expr.l_or(lhs, rhs)` ¶

Logical or of lhs and rhs

NOTE: Short-circuit evaluation is NOT supported

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.l_or

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The logical or

`freetensor.core.expr.le(lhs, rhs)` ¶

lhs <= rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.le

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The comparison

`freetensor.core.expr.ln(expr)` ¶

Natural logarithm

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.ln

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The exponent

`freetensor.core.expr.load_at_version(tape_name, dtype, *indices)` ¶

Create an LoadAtVersion node (only for custom gradient)

This node is only used for custom gradient. See UserGradForPrevStmt.

`freetensor.core.expr.lt(lhs, rhs)` ¶

lhs < rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.lt

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The comparison

`freetensor.core.expr.max(lhs, rhs)` ¶

Maximum of lhs and rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.max

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The maximum

`freetensor.core.expr.min(lhs, rhs)` ¶

Minimum of lhs and rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.min

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The minimum

`freetensor.core.expr.mod(lhs, rhs)` ¶

lhs modulus rhs

This function follows floored division (following Python convention, instead of C). See https://en.wikipedia.org/wiki/Modulo for details. This function is recommended over remainder, as it enjoys more optimizations

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.mod

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The modulo

`freetensor.core.expr.mtype(var)` ¶

Get memory type of a variable

`freetensor.core.expr.mul(lhs, rhs)` ¶

lhs * rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.mul

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The product

`freetensor.core.expr.ndim(var)` ¶

Get the number of dimensions of a variable

`freetensor.core.expr.ne(lhs, rhs)` ¶

lhs != rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.ne

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The comparison

`freetensor.core.expr.neg(expr)` ¶

Negation

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.neg

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The negation

`freetensor.core.expr.remainder(lhs, rhs)` ¶

Remainder of lhs dividing rhs

This function follows truncated division (following C convention, instead of Python). See https://en.wikipedia.org/wiki/Modulo for details. End users are encouraged to use lhs % rhs instead, which follows Python convetion, and enjoys better optimization in FreeTensor

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.remainder

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The remainder

`freetensor.core.expr.round_towards_0_div(lhs, rhs)` ¶

C-style integer division of lhs dividing by rhs

The result rounds towards 0 (following C convention, instead of Python) End users are encouraged to use lhs // rhs instead, which follows Python convetion, and enjoys better optimization in FreeTensor

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.round_towards_0_div

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The quotient

`freetensor.core.expr.shape(var, i=None)` ¶

shape(var, i): Get size of specified dimension of a variable

shape(var): Get sizes of all dimensions of a variable

`freetensor.core.expr.sigmoid(expr)` ¶

Sigmoid

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.sigmoid

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.sin(expr)` ¶

Sine

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.tanh

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.sqrt(expr)` ¶

Square root

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.sqrt

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The square root

`freetensor.core.expr.square(expr)` ¶

Square

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.square

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The square

`freetensor.core.expr.sub(lhs, rhs)` ¶

lhs - rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.sub

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The difference

`freetensor.core.expr.tan(expr)` ¶

Tangent

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.tanh

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.tanh(expr)` ¶

Hyperbolic tangent

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.tanh

Parameters:	expr (`VarRef or Number`) – The operand

Returns:	`VarRef or Number` – The result

`freetensor.core.expr.truediv(lhs, rhs)` ¶

Floating point division of lhs dividing by rhs

For scalar operands, it emit an expression node in AST. For non-scalar operands, it calls libop.truediv

Parameters:	lhs (`VarRef or Number`) – The left-hand-side operand rhs (`VarRef or Number`) – The right-hand-side operand

Returns:	`VarRef or Number` – The quotient

`freetensor.core.expr.unbound(expr)` ¶

Explicitly mark a bool expression not to used for boundary analysis and thus

not used for shrinking loops or variables

This is useful when we want to leave a loop length constatnt or regular, and guard it with a branch, e.g.:

for i = 0 to 100
  if unbound(i < x)
    ...

Parameters:	expr (`VarRef or Number`) – The bool expression to mark

Returns:	`VarRef or Number` – The marked expression

`freetensor.core.frontend` ¶

A frontend transforming user Python functions to ASTs via staging.

`freetensor.core.frontend.FreeTensorOverload (StagingOverload)` ¶

Helper class managing context in IR staging.

`freetensor.core.frontend.FreeTensorOverload.allow_shortcut_scope(self, allow)` `inherited` ¶

Opens a scope that allows shortcut control flows in a statically deterministic

context. Need to be closed by with statement.

`freetensor.core.frontend.FreeTensorOverload.assert_stmt(self, test)` `inherited` ¶

Assert staging tool.

`freetensor.core.frontend.FreeTensorOverload.assign_stmt(self, name, value)` `inherited` ¶

Customized assign wrapper.

If value is instance of StagedAssignable, it's regarded as a customized assign behavior and gets executed with the assigned target variable name. This wrapper is used for initializing a variable.

`freetensor.core.frontend.FreeTensorOverload.break_stmt(self)` `inherited` ¶

Break staging tool. Only allow break in static control flow.

`freetensor.core.frontend.FreeTensorOverload.continue_stmt(self)` `inherited` ¶

Continue staging tool. Only allow continue in static control flow.

`freetensor.core.frontend.FreeTensorOverload.custom_attr(self, obj, attr)` ¶

Customized attribute accessor.

The framework first looks for a Python native attribute. If not found, it looks for this overloaded custom attribute resolver.

The default implementation provides no custom attribute. Can be overridden by subclasses.

Parameters:	obj (`Any`) – Object to access attribute. attr (`str`) – Attribute name.

Returns:	`Any` – The attribute value.

`freetensor.core.frontend.FreeTensorOverload.foreach(self, names, iter, body)` `inherited` ¶

Customized foreach wrapper.

If value is instance of StagedIterable, its regarded as a customized foreach behavior and used to generate code for the python for loop. Otherwise, we try to execute the loop as usual.

`freetensor.core.frontend.FreeTensorOverload.fullname(self, name)` ¶

Get distinct name.

`freetensor.core.frontend.FreeTensorOverload.functiondef_wrapper(self, filename, func)` ¶

Function definition wrapper.

This wrapper performs extra initialization and cleanup for function definition.

`freetensor.core.frontend.FreeTensorOverload.if_then_else_expr(self, predicate, then_expr, else_expr)` `inherited` ¶

If-then-else expression staging tool.

`freetensor.core.frontend.FreeTensorOverload.if_then_else_stmt(self, predicate, then_body, else_body=None)` `inherited` ¶

If-then-else statement staging tool.

When predicate is deterministic in staging, only one branch is generated. Otherwise, a If node in IR is generated.

`freetensor.core.frontend.FreeTensorOverload.load_attr(self, obj, attr)` `inherited` ¶

Load attribute staging tool. Allows customization of reading attributes.

`freetensor.core.frontend.FreeTensorOverload.mark_position(self, lineno)` ¶

Code position handler.

Defaults to a no-op. Can be overridden by subclasses.

Parameters:	lineno (`int`) – Line number of the next statement.

`freetensor.core.frontend.FreeTensorOverload.metadata(self, entry)` ¶

Metadata handler.

A metadata line is a comment starting with #! and followed by a metadata, represented as a string parameter.

Defaults to a no-op. Can be overridden by subclasses.

Parameters:	content (`str`) – The metadata content.

`freetensor.core.frontend.FreeTensorOverload.return_stmt(self, value, funcname)` `inherited` ¶

Return staging tool. Only allow return in static control flow.

`freetensor.core.frontend.FreeTensorOverload.unpack_assign_stmt(self, names, values)` `inherited` ¶

Customized assign wrapper for one or more targets.

If values is instance of StagedUnpackAssignable, it's regarded as a customized assign behavior and gets executed with all the assigned targets' names. Otherwise, it calls assign_stmt with each sub-assignments.

Please note that names can be nested tuples like ("a", ("b", "c")).

Please also note that names can also be a single string like "a" even if values is a tuple. There is no unpacking in this case

`freetensor.core.frontend.FreeTensorOverload.while_stmt(self, fpred, body)` `inherited` ¶

While statement staging tool.

`freetensor.core.frontend.FreeTensorStagingError (StagingError)` ¶

Error occurred during staging function execution (i.e. IR tree generation).

`freetensor.core.frontend.LifetimeScope` ¶

This scope is used to register multiple scopes inside a single lifetime scope.

The inner scopes might be used to register variables, etc. They will be exited in reverse order of their registration.

`freetensor.core.frontend.PredefinedVarCreator (VarCreator)` `dataclass` ¶

`freetensor.core.frontend.PredefinedVarCreator.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.frontend.PredefinedVarCreator.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.frontend.PredefinedVarCreator.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.frontend.PredefinedVarCreator.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.frontend.PredefinedVarCreator.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.PredefinedVarCreator.init(self, initializer, dtype, mtype)` `special` ¶

Initialize self. See help(type(self)) for accurate signature.

`freetensor.core.frontend.PredefinedVarCreator.assign(self, name)` ¶

Customized assign behavior. Creates a VarDef with its full name.

`freetensor.core.frontend.ReversedIterable (StagedIterable)` ¶

`freetensor.core.frontend.ReversedIterable.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.frontend.ReversedIterable.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.frontend.ReversedIterable.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.frontend.ReversedIterable.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.frontend.ReversedIterable.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.UserGrad (UserGradStaged)` ¶

Define a custom gradient

Follow the following steps to define custom gradient:

Add some mark_version statements in the program. mark_version('y0', y) marks the specific versions of variable y at the program position of the statement and at all iterations as 'y0'.
Add a UserGrad scope. 2.1. UserGrad optionally receives parameter stmt_range, recorded by the StmtRange helper class, which means the gradient is for the code specified in the range. Ignoring the parameter means setting gradient for the previous statement of the scope. 2.2. Other parameters of UserGrad sets the mapping from original variables to gradient variables. with UserGradForPrevStmt(x, y) as (dx, dy) provides VarRef dx and dy as gradient variables to be used inside the scope.
In order to use the value from the forward pass in the backward pass, do not access the forward variables directly in the scope. Instead, use load_at_version expressions. load_at_version(y0, i, j) loads from y[i, j] at the specific version marked by y0 = mark_version(y), saved from the same iteration in the forward pass. (If directly writing staged code, it is MarkVersion('y0', y)). In other words, after AD, the position of mark_version and the dynamic loop iterator together makes up the actual version number for the tape.
Build the AST with pop_ast_and_user_grads instead of pop_ast. An extra list will be returned together with the AST, which you need to pass as grad's user_grads argument. This list records the forward-to-backward relation of the nodes.

If you are directly writing staged code, use UserGradStaged instead.

Parameters:	args (`Sequence[freetensor.core.expr.VarRef]`) – (Positional variadic) Mapping from original variables to gradient variables stmt_range* (`StmtRange`) – The range in the original program that we are setting custom gradient for

`freetensor.core.frontend.Var (StagedTypeAnnotation)` ¶

`freetensor.core.frontend.Var.class (ABCMeta)` `inherited` ¶

`freetensor.core.frontend.Var.class.base (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

freetensor.core.frontend.Var.__class__.__base__.__instancecheck__(cls, instance) special ¶

Override for isinstance(instance, cls).

freetensor.core.frontend.Var.__class__.__base__.__new__(mcls, name, bases, namespace, **kwargs) special staticmethod ¶

Create and return a new object. See help(type) for accurate signature.

freetensor.core.frontend.Var.__class__.__base__.__subclasscheck__(cls, subclass) special ¶

Override for issubclass(subclass, cls).

freetensor.core.frontend.Var.__class__.__base__.register(cls, subclass) ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.Var.class.register(cls, subclass)` `inherited` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.Var.init(self, shape, dtype, atype='input', mtype=None)` `special` ¶

Declare a variable

Parameters:

shape (Sequence[Expr] or Var) – Shape of the variable. A variable can be created using a literal shape, or another fixed-length VarRef as a shape
dtype (str or DataType) – Data type of the variable
atype (str or AccessType) – Access type of the variable. It specifies whether (and how) the variable is an I/O variable of the function it belongs to. Defaults to "input"
mtype (str or MemType (Optional)) – Memory type of the variable. If omitted, the main memory type of the default Target in config will be used

`freetensor.core.frontend.VarCreator (StagedAssignable)` `dataclass` ¶

VarCreator(shape: Union[Sequence, freetensor.core.expr.VarRef], dtype: str, mtype: str, assigned: bool = False)

`freetensor.core.frontend.VarCreator.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.frontend.VarCreator.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.frontend.VarCreator.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.frontend.VarCreator.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.frontend.VarCreator.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.VarCreator.assign(self, name)` ¶

Customized assign behavior. Creates a VarDef with its full name.

`freetensor.core.frontend.VersionMarker (StagedAssignable)` `dataclass` ¶

VersionMarker(var: freetensor.core.expr.VarRef, assigned: bool = False)

`freetensor.core.frontend.VersionMarker.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.frontend.VersionMarker.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.frontend.VersionMarker.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.frontend.VersionMarker.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.frontend.VersionMarker.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.VersionMarker.assign(self, tape_name)` ¶

Customized assign behavior. Creates a MarkVersion with its full name.

`freetensor.core.frontend.dynamic_range (StagedIterable)` ¶

Dynamic range that generates For loop in IR tree.

`freetensor.core.frontend.dynamic_range.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.frontend.dynamic_range.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.frontend.dynamic_range.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.frontend.dynamic_range.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.frontend.dynamic_range.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.frontend.dynamic_range.init(self, start, stop=None, step=1)` `special` ¶

Initialize a dynamic range.

Arguments semantic identical to builtin range.

`freetensor.core.frontend.capture_var(*args, **kwargs)` ¶

Capture external array as tensor variable.

`freetensor.core.frontend.empty(*args, **kwargs)` ¶

Create an empty variable

Parameters:	shape (`Sequence[Expr] or Var`) – Shape of the variable. A variable can be created using a literal shape, or another fixed-length VarRef as a shape dtype (`str or DataType`) – Data type of the variable mtype (`str or MemType (Optional)`) – Memory type of the variable. If omitted, the main memory type of the default Target in config will be used

`freetensor.core.frontend.push_for_backward(var)` ¶

Push the current value from the forward pass to be used at the backward pass

This function is for custom gradients. See UserGrad for details on how to provide custom gradients.

You may imagine there is a virtual stack for each variable. Each time you call x_handle = push_for_backward(x) in the forward pass, the value of x at the current iteration will be "pushed" to the virtual stack. You can access x_handle at the backward pass. Each time you access x_handle, you will "pop" the stack and get the value of x pushed at the same iteration. Since the "stack" is virtual, you do NOT need to "pop" the same count as "push"es: the version numbering is fully automatic. Besides, there may not be a real stack at runtime: it can be compiled to any data structure.

This function will be staged to mark_version statement in the IR.

`freetensor.core.frontend.var(*args, **kwargs)` ¶

Create an with variable a given initializer

Parameters:

initializer (Sequence[Sequence[...Sequence[Expr]...]]) – (Multi-level of) sequence of expressions. Will be data of the variable
shape (Sequence[Expr] or Var) – Shape of the variable. A variable can be created using a literal shape, or another fixed-length VarRef as a shape
dtype (str or DataType) – Data type of the variable
mtype (str or MemType (Optional)) – Memory type of the variable. If omitted, the main memory type of the default Target in config will be used

`freetensor.core.func` ¶

`freetensor.core.func.Func (EnableAttachBackward, Func)` ¶

`freetensor.core.func.Func.call(self, *args, **kvs)` `special` ¶

Enable invoking a transformed AST in another function being transformed, via

inlined_invoke

`freetensor.core.jit` ¶

`freetensor.core.jit.JIT` ¶

Declare a function parameter as a JIT parameter

A function with one or more JIT parameters will be compiled to a JIT template. It can be instantiate after the JIT paraemters are provided

Usage: x: JIT or x: JIT[AnyPythonType]. The latter form has no syntactic meanings, and is only for documentation.

NOTE 1: The JIT type annotation can only be used for parameter of the outer-most function intended for @transform (or @optimize, etc). It can NOT be used for inner functions intended for @inline.

NOTE 2: The JIT type annoation can only annotated on parameters inside the function signature. It is NOT supported in annotations for statements.

`freetensor.core.jit.JITTemplate (ABC)` ¶

A template that can be instantiated given concrete arguments

By calling instantiate with actual arguments you are expecting to run a JIT function with, an instantiated object will be returned. Subclasses of JITTemplate should override instantiate_by_only_jit_args, and define what is actually returned.

Parameters:	params (`OrderedDict`) – Parameter list from inspect.signature(func).parameters jit_param_names (`Sequence`) – Sequence of names of JIT parameters in the original order defined in the function

`freetensor.core.jit.JITTemplate.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.jit.JITTemplate.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.jit.JITTemplate.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.jit.JITTemplate.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.jit.JITTemplate.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.jit.JITTemplate.instantiate(self, *args, **kvs)` ¶

Get an instance with the arguments you are expecting to run a JIT function with

`freetensor.core.jit.JITTemplate.instantiate_and_call(self, *args, **kvs)` ¶

Get an instance and call it with the arguments you are expecting to run a JIT function with

`freetensor.core.jit.JITTemplate.instantiate_by_only_jit_args(self, *jit_args)` ¶

Get an instance with only JIT arguments

This function accpets a tuple of arguments. Keyword arguments is NOT supported, so memoization can be easier, with considering the order of the arguments.

`freetensor.core.jit.JITTemplate.separate_args(self, *args, **kvs)` ¶

Return a list of non-JIT args, and a list of JIT args

`freetensor.core.optimize` ¶

`freetensor.core.optimize.optimize(func=None, schedule_callback=None, backward_schedule_callback=None, target=None, device=None, default_dynamic_range=True, jit_cache=<function cache at 0x147b6469ca60>, verbose=0)` ¶

An one-click optimization from Python function to binary executable

Usage:

@optimize
def f(...):
    ...

It is equivalent to:

@build_binary
@codegen
@lower
@transform
def f(...):
    ...

Parameters:

func (Python function or AST) – The user function to optimize. If not specified, a partial function will be returend, which can be used as a decorator
schedule_callback (Optional[Callable[[freetensor.core.schedule.Schedule], NoneType]]) – Schedule(s) to apply
backward_schedule_callback (Callable[[freetensor.core.schedule.Schedule], NoneType]) – Specify what schedule(s) to do for the backward function, if ast is returned from AD with attach_backward=True. Defaults to be the same with callback
target (Optional[freetensor_ffi.Target]) – The target architecture. You don't have to set target if you set device
device (Optional[freetensor_ffi.Device]) – Where to run the program
default_dynamic_range (bool) – If True, the built-in range is replaced with freetensor.dynamic_range. Defaults to True
verbose (int) – Verbosity level. Can be 0, 1 or 2

`freetensor.core.optimize.optimize_to_pytorch(func=None, tapes=<GradTapeMode.NoReuseOnly: 2>, forward_schedule_callback=None, backward_schedule_callback=None, target=None, device=None, default_dynamic_range=True, verbose=0)` ¶

Compile a FreeTensor function to a PyTorch call, whose gradient can be

recognized by PyTorch

The compiled function will be a typical PyTorch's "function" (rather than a PyTorch's "module"). Technically, this means it is a wrapper function around a PyTorch's Function's apply method

Schedules (if any) must be applied to the forward function and the backward function separated. For this reason, currently only first-order gradient is supported

Parameters:

func (Python function or AST) – The user function to optimize. If not specified, a partial function will be returend, which can be used as a decorator
tapes (Union[Sequence, freetensor_ffi.GradTapeMode]) – Intermediate variables that need to be stored from the forward pass and reused in the backward pass. This parameter can be a sequence, which contains VarDef selectors of them. It can also be a GradTapeMode, then it will determine which intermediate variables to be stored by heuristics. Avail GradTapeModes are: All: store all variables including local scalars; None: store nothing; NoReuseOnly: store variables that only hold one version of data, which means we do not have to store each version of them in their history
forward_schedule_callback (Optional[Callable[[freetensor.core.schedule.Schedule], NoneType]]) – Schedule(s) to apply to the forward function
backward_schedule_callback (Optional[Callable[[freetensor.core.schedule.Schedule], NoneType]]) – Schedule(s) to apply to the backward function
target (Optional[freetensor_ffi.Target]) – The target architecture. You don't have to set target if you set device
device (Optional[freetensor_ffi.Device]) – Where to run the program
default_dynamic_range (bool) – If True, the built-in range is replaced with freetensor.dynamic_range. Defaults to True
verbose (int) – Verbosity level. Can be 0, 1 or 2

`freetensor.core.param_ret_dict` ¶

`freetensor.core.param_ret_dict.ParamRetDict` ¶

Look an object using either a function parameter or return value's name or position

`freetensor.core.param_ret_dict.ParamRetDict.init(self, d, *, func=None, func_name=None, param_names=None, return_names=None)` `special` ¶

Either func or (func_name and param_names and return_names) should be provided

`freetensor.core.param_ret_dict.Parameter` ¶

Alias of a parameter of a function by position instead of by name

Parameter(n) represents the n-th parameter (counted from 0)

Parameter() can be used if there is only one parameter

`freetensor.core.param_ret_dict.Return` ¶

Alias of a return value of a function by position instead of by name

Return(n) represents the n-th return value (counted from 0)

Return() can be used if there is only one return value

`freetensor.core.passes` ¶

`freetensor.core.passes.lower(ast=None, target=None, skip_passes=[], jit_cache=<function cache at 0x147b6469ca60>, verbose=0)` ¶

Lower an AST using a series of passes

Parameters:

ast (AST) – The AST to be lowered. Can be a Func or a Stmt. If not specified, a partial function of lower will be returned, which can be used as a decorator
target (Optional[freetensor_ffi.Target]) – Lower the AST to a target with target-specific passes, then the AST can be used for codegen. If not set, use the default Target in Config
skip_passes (Sequence[str]) – Skip some pass for testing or debugging. Names in skip_passes are in underscore_style, as in Python. Please note that some passes will not be skipped even specified in these parameter, because they are indirectly called in some other passes
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
verbose (int) – 0 = print nothing. 1 = print the lowered AST. 2 = print AST after every single passes

Returns:	`Func or JITTemplate` – Return a Func for an AST if there is no JIT parameters. Return a JITTemplate that generates a Func if there is at least one

`freetensor.core.return_values_pack` ¶

`freetensor.core.return_values_pack.ReturnValuesPack` ¶

Hold return values from a Driver invocation

Return values can be retrieved in an anonymous manner: x, y, z = pack, or in a named manner: pack['x']

Please note that a ReturnValuesPack is different from a OrderedDict, as OrderedDict unpacks to keys rather than values

`freetensor.core.return_values_pack.ReturnValuesPack.contains(self, key)` `special` ¶

Test if a return value exists

`freetensor.core.return_values_pack.ReturnValuesPack.getitem(self, key)` `special` ¶

Get a return value with a name. Tuple is supported for multiple values

`freetensor.core.return_values_pack.ReturnValuesPack.iter(self)` `special` ¶

Get all return values in the order declared in Func

`freetensor.core.schedule` ¶

`freetensor.core.schedule.IDMap` ¶

A dict-like container recording an ID-to-ID mapping, representing what IDs

become what IDs after a schedule

An IDMap can be looked up by numerical ID, or by Stmt instances or Selector strings of the original (before applying schedule) AST

`freetensor.core.schedule.Schedule (Schedule)` ¶

`freetensor.core.schedule.Schedule.as_matmul(self, loop, mode=<AsMatMulMode.KeepMemLayout: 0>, target=None, backend=None)` ¶

Transform nested loops to be a external call to a matrix multiplication

Parameters:

loop (str, ID or Stmt) – ID of the loop
allow_var_reorder (bool) – If true, automatically try calling var_reorder to eanble as_matmul
mode (AsMatMulMode) – What to do if the memory layout does not meet the requirement from the external library. KeepMemLayout => Raise an exception. TryVarReorder => try var_reorder on some variables, but may affect performance of other use of these variable. TryTranspose => try cache and then var_reorder on some variables, but will incur extra overhead.
target (Target) – Hardware target. If omitted, use the default target in config, or the target set by with scopes.
backend (Union[str, freetensor_ffi.MatMulBackend]) – Backend library. Defaults to "mkl" for CPU targets, "cublas" for GPU targets.

Exceptions:	`InvalidSchedule` – if the loop cannot be transformed to be a matrix multiplication

`freetensor.core.schedule.Schedule.ast(self)` ¶

Get the scheduled AST without function signature

This is mainly for debugging and testting purpose

`freetensor.core.schedule.Schedule.auto_fission_fuse(self, target)` ¶

(Experimental) Automatically fuse consecutive loops or vice versa using

some heuristics

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_inline(self, target)` ¶

(Experimental) Automatically inline very-small VarDef nodes

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_mem_layout(self, target)` ¶

(Experimental) Automatically adjust memory layout of variables

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_parallelize(self, target)` ¶

(Experimental) Automatically parallelize some loops using some heuristics

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_pluto(self, target)` ¶

(Experimental) Automatically apply pluto-based schedules

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_schedule(self, target)` ¶

(Experimental) Automatic scheduling using some heuristics

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_set_mem_type(self, target)` ¶

(Experimental) Automatically set memory types using some heuristics

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_swap(self, target)` ¶

(Experimental) Automatically swap statements to enable more fission or

fusion

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_unroll(self, target)` ¶

(Experimental) Automatically unroll loops using some heuristics

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.auto_use_lib(self, target)` ¶

(Experimental) Automatically use external libs using some heuristics

Parameters:	target (`Target`) – Target architecture

`freetensor.core.schedule.Schedule.blend(self, loop)` ¶

Unroll a loop and interleave statements from each iteration

E.g.

for i = 0 to 2 {
f(i);
g(i);
}

will be transformed to be

f(0);
f(1);
g(0);
g(1);

Virtual threads in TVM can be implemented via blend

Parameters:	loop (`str, ID or Stmt`) – The loop being transformed

Exceptions:	`InvalidSchedule` – if the loop is not found, the loop length is not a constant, or the dependences cannot be solved

`freetensor.core.schedule.Schedule.cache(self, stmt, var, mtype)` ¶

Cache a variable into a new local variable

All needed data will be filled into the cache first, then all reads and writes will be directed to the cache, and finally all needed data will be flushed from the cache

Note for reduction: This transformation preserves the computation order. It will transform

a += x
a += y

to

a.cache = a + x + y
a = a.cache

If you need a "real" cache for reduction, which reorders the computation, use cache_reduction instead

Parameters:	stmt (`str, ID or Stmt`) – The statement or block (e.g. an If or a For) to be modified var (`str`) – Name of the variable to be cached mtype (`MemType`) – Where to cache

Exceptions:	`InvalidSchedule` – if the ID or name is not found

Returns:	`(ID, ID, ID, ID)` – (ID of the statement that fills the cache, ID of the statement that flushes from the cache, name of the cache variable, ID of the VarDef node of the cache variable)

`freetensor.core.schedule.Schedule.cache_reduction(self, stmt, var, mtype)` ¶

Perform local reductions (e.g. sum) in a local variable first, and then

reduce the local result to the global variable

E.g.

a += x
a += y

will be transformed to be

a.cache = x + y
a += a.cache

Parameters:	stmt (`str, ID or Stmt`) – The statement or block (e.g. an If or a For) to be modified var (`str`) – Name of the variable to be cached. Only reductions are allowed on `var` in `stmt`. Plain reads or writes are not allowed mtype (`MemType`) – Where to cache

Exceptions:	`InvalidSchedule` – if the ID or name is not found, or there are unsupported reads or writes

Returns:	`(ID, ID, ID, ID)` – (ID of the statement that initialize the cache, ID of the statement that reduces the local result to the global result, name of the cache variable, ID of the VarDef node of the cache variable)

`freetensor.core.schedule.Schedule.fission(self, loop, side, splitter, allow_enlarge=True)` ¶

Fission a loop into two loops each containing part of the statements, one

followed by another

To split loop into two nested loops, use split instead

Statements inside the original loop will be distributed to one or both (happening if they are scope statements) loops. If a statement is originally labeled "S", it can be selected by "$fission.0{S}" (from the first loop) or "$fission.1{S}" (from the second loop) after fission. If one of the resulting loop has an empty body, it will be removed

Parameters:

loop (str, ID or Stmt) – The loop to be fissioned
side (FissionSide) – If After, splitter is the last statement of the first loop. If Before, splitter is the first statement of the second loop
splitter (str (Selector string), ID, Stmt, or list of them) – Where to fission the loop. If multiple statement are selected, fission the look before or after all of them
allow_enlarge (bool) – If True, try to avoid dependence by enlarging some VarDef nodes. If False, raise InvalidSchedule in such cases.

Exceptions:	`InvalidSchedule` – if any dependence cannot be resolved

Returns:	`(IDMap, IDMap)` – ({old ID -> new ID in 1st loop}, {old ID -> new ID in 2nd loop}). If a loop is removed because it has an empty body, it will not be in the returned map

`freetensor.core.schedule.Schedule.fork(self)` ¶

fork(self: freetensor_ffi.Schedule) -> freetensor_ffi.Schedule

`freetensor.core.schedule.Schedule.func(self)` ¶

Get the scheduled function

`freetensor.core.schedule.Schedule.fuse(self, loop0, loop1=None, strict=False)` ¶

Fuse two directly following loops with the same length into one

To merge nested loops into one, use merge instead

parallelize, unroll and vectorize properties will be reset on the fused loop

Suppose the original loops are labeled "L1" and "L2", the fused loop can be selected by "$fuse{L1, L2}"

Parameters:	loop0 (`str, ID or Stmt`) – The leading loop loop1 (`str, ID or Stmt, Optional`) – The following loop. If omitted, it will try to find a following loop of `loop0` strict (`bool`) – False by default. If set to True, throw an error if unable to determine whether the two loops are of the same length

Exceptions:	`InvalidSchedule` – if the two loops are not directly following, the two loops are not of the same length, or there is any dependence cannot be resolved

Returns:	`ID` – ID of the result loop

`freetensor.core.schedule.Schedule.inline(self, vardef)` ¶

Remove a variable. When the variable is used, recompute its value

Parameters:	vardef (`str, ID or Stmt`) – The VarDef statement of the specific variable. It can not be an I/O varible

Exceptions:	`InvalidSchedule` – if the variable cannot be completely removed

`freetensor.core.schedule.Schedule.merge(self, loop1, loop2)` ¶

Merge two directly nested loops into one

To fuse consecutive loops, use fuse instead

parallelize, unroll and vectorize properties will be reset on the merged loop

Suppose the original loops are labeled "L1" and "L2", the merged loop can be selected by "$merge{L1, L2}"

Parameters:	loop1, loop2 (`str, ID or Stmt`) – loops to be merged, can be in any order

Exceptions:	`InvalidSchedule` – if the loops are not directly nested

Returns:	`ID` – ID of the merged loop

`freetensor.core.schedule.Schedule.move_to(self, stmt, side, dst)` ¶

Move a statement to a new position

This is a composite schedule command, which is implemented with other commands

If moving a statement out of some loops, identical loops will be added around the moved statement, which is equivalent to fission these loops

Parameters:	stmt (`str, ID or Stmt`) – The statement to be moved side (`MoveToSide`) – Whether `stmt` will be BEFORE or AFTER `dst dst (`str (Selector string), ID, Stmt, or list of them`) – Insert `stmt` to be directly after this statement. If multiple statements are selected, move to before or after all of them

Exceptions:	`InvalidSchedule` – if there is no feasible path to move

Returns:	`(ID, ID)` – (The new ID of the moved statement, The out-most newly introduced statments including the added loops)

`freetensor.core.schedule.Schedule.parallelize(self, loop, parallel)` ¶

Mark a loop with a parallel implementation

This schedule follows a fork-join model: multiple workers (abstract threads) are created (but physically the threads may be cached in a thread pool) when the loop begins, do their jobs in parallel, and join when the loop ends

OpenMP threads follow a typical fork-join model. CUDA threads run in a bulk-synchronous parallel (BSP) model, which can also be mimiked by the fork-join model: All threads start when the kernel get launched, but they only begin to do their jobs when the parallel loop begins. Nevertheless, the fork-join model needs the following extension to fully mimic a BSP model:

Taking CUDA as an example, we allow binding a loop to threadIdx.x inside another loop bound to threadIdx.x, which is illegal in a classic fork-join model. For example, we may implement a matmul with collaborative fetch as below:

for i : threadIdx.x  # Li
  for j : threadIdx.y  # Lj
    local_sum = 0  # In gpu/local memory, unique to (i, j)
    for k0  # Lk0
      for k : threadIdx.y  # Lk1_a
        A_cache[k] = A[i, k]  # In gpu/shared, shared by different j
      for k : threadIdx.x  # Lk1_b
        B_cache[k] = B[k, j]  # In gpu/shared, shared by different i
      for k  # Lk1_c
        sum += A_cache[k] * B_cache[k]
    C[i, j] = local_sum

A seemingly plausible solution to avoid this extension is to reorder Lk0 to outer-most, and then move Lk1_a and Lk1_b out of Li or Lj. This resolves the nested threadIdx.x and threadIdx.y binding problem by running Li+Lk1_a, Lj+Lk1_b and Li+Lj interleavingly, instead of running Lk1_a and Lk1_b inside Li+Lj. However, this approach is illegal, because the local variable local_sum can no longer be kept inside the body of Li and Lj: It has to be reused across multiple runs of Li and Lj

Please also note that we can bind one threadIdx.x to two loops only when the body statement is loop-invariant to one of them. For example, the following binding is still illegal, even in our extended fork-join model, because it violates its serial semantics:

for i : threadIdx.x
  for j : threadIdx.x
    A[i, j] ++

Parameters:	loop (`str, ID or Stmt`) – The loop parallel (`ParallelScope`) – Parallel scope allow_reduction (`bool`) – If false, raise InvalidSchedule if this schedule would introduce a parallel reduction

Exceptions:	`InvalidSchedule` – if the loop is not found or unable to be parallelized

`freetensor.core.schedule.Schedule.parallelize_as(self, nest, reference, vardef)` ¶

Parallelize a loop nest according to another loop nest to keep a tensor

thread-local

Parameters:

nest (str, ID or Stmt) – The loop nest to be parallelized. The ID can be of any statement type, and all statements it contains will be parallelized.
reference (str, ID or Stmt) – The loop nest to be referenced. The ID can be of any statement type, and all statements it contains will be referenced.
vardef (str, ID or Stmt) – The VarDef statement of the tensor to be kept thread-local.

Exceptions:	`InvalidSchedule` – if any of the ID is not found, or the reference loop nest is already thread-non-local.

`freetensor.core.schedule.Schedule.permute(self, loops, transform_func)` ¶

Permute perfectly nested loops (directly nested loops without statements

in between) with the given loop space transformation function

The transformed loops follow ascending lexical order of the transformed terms returned by transformFunc when called with original iteration

Parameters:	loops (`array like of str, ID or Stmt`) – the list of perfectly nested loops to be permuted transform_func (`Callable[[Expr], Expr]`) – the loop space transformation function, should be bijective

Returns:	`list of ID` – the list of IDs of permuted loops

`freetensor.core.schedule.Schedule.pluto_fuse(self, loop0, loop1, nest_level_0=0, nest_level_1=0, fusable_overlap_threshold=1, fusable_nonoverlap_tolerance=4, do_simplify=True)` ¶

Use Pluto+ algorithm to permute and fuse two loops, with as most parallelizable

loops as possible at outermost levels. The two loops are required to be consequent; all directly nested levels are detected and subject to permutation. Remaining levels that cannot be fused are left inside the fused loops as two statements

Parameters:

loop0 (str, ID or Stmt) – The first loop to fuse
loop1 (str, ID or Stmt) – The second loop to fuse
nest_level_0 (int) – The number of nesting levels of loop 0 to be considered, defaults to maximum possible
nest_level_1 (int) – The number of nesting levels of loop 1 to be considered, defaults to maximum possible
fusable_overlap_threshold (int) – The minimum overlapping size of two loops to be regarded fusable. Defaults to 1
fusable_nonoverlap_tolerance (int) – The maximum non-overlapping size at either side of two loops to be regarded fusable. Defaults to 4
do_simplify (bool) – Whether the result is simplified by the way, defaults to true

Exceptions:	`InvalidSchedule` – if the loops are not consequent

Returns:	`(ID, int)` – The ID of fused loop and level of parallelizable loops

`freetensor.core.schedule.Schedule.pluto_permute(self, loop, nest_level=0, do_simplify=True)` ¶

Use Pluto+ algorithm to permute a single loop, with as most parallelizable loops

as possible at outermost levels.

Parameters:	loop (`str, ID or Stmt`) – The loop to permute nest_level (`int`) – The number of nesting levels to be considered, defaults to maximum possible do_simplify (`bool`) – Whether the result is simplified by the way, defaults to true

Returns:	`(ID, int)` – The ID of permuted loop and level of parallelizable loops

`freetensor.core.schedule.Schedule.reorder(self, order, mode=<ReorderMode.PerfectOnly: 0>)` ¶

Reorder directly nested loops

To swap consecutive loops, use swap instead

Parameters:

order (array like of str, ID or Stmt) – Vector of loops. The requested order of the loops
mode (ReorderMode) – How to deal with imperfectly nested loops. PerfectOnly => raise an exception. MoveOutImperfect => do fission in advance to move out statements between the loops, which may enlarge intermediate tensors. MoveInImperfect => move statements between the loops inwards after adding gurads them them, which may hurt parallelism

Exceptions:	`InvalidSchedule` – if the input is invalid or there are breaking dependences

`freetensor.core.schedule.Schedule.separate_tail(self, noDuplicateVarDefs=False)` ¶

Seperate main iterations and tail iterations of a loop

E.g.

for i = 0 -> 3 {
  for j = 0 -> 4 {
     if (i * 4 + j < 10) {
       ...
     }
  }
}

Each loop will be separated into 2 parts: the body and the tail. After simplification, the program will finally be transformed to

for i = 0 -> 2 {
  for j = 0 -> 4 {
    ...
  }
}
for j = 0 -> 2 {
  ...
}

Ideally, all programs can benefit from this schedule. However, this schedule may greatly increase the program size and make the compiling time way too long. Therefore, this transformation is implemented as a schedule, which can be applied optionally. (TODO: Optionally apply this schedule to part of the program)

Parameters:	noDuplicateVarDefs (`bool`) – If there is two VarDef nodes in two branches, it may result in doubled memory use, since different thread may go to different branch. Set this parameter to true to stop duplicating VarDef nodes.

`freetensor.core.schedule.Schedule.set_mem_type(self, vardef, mtype)` ¶

Change where a variable is stored

Parameters:

vardef (str, ID or Stmt) – ID of the VarDef statement of the specific variable
mtype (MemType) – Where the variable should be stored
rejectIndirectAccess (bool) – Registers usually do not support indirect access. If a variable is accessed indirectly, setting it to use registers is meaningless even successful. If this parameter is set to true, throw an exception if the variable being set is accessed indirectly. Specifically, two types of access are considered indirect: 1) The index is a load from another variable, or 2) The index is a loop iterator and the loop has a dynamic length (which can not be unrolled by a backend compiler). By default, this parameter is determined automatically by mtype.

Exceptions:	`InvalidSchedule` – if the variable is not found, or if rejecting an indirect access

`freetensor.core.schedule.Schedule.split(self, node, factor=-1, nparts=-1, shift=0)` ¶

Split a loop into two nested loops

To fission a loop into two consecutive loops, use fission instead

Two modes are provided:

Specify factor and leave nparts to -1. It will result in an outer loop with length ceil(n / factor), and an inner loop with length factor, where n is the original loop length added by shift. The original iterator i will be transformed to i0 * factor + i1, where i0 and i1 are the iterators of the new outer and inner loops, respectively
Specify nparts and leave factor to -1. It will result in an outer loop with length nparts, and an inner loop with length ceil(n / nparts), where n is the original loop length added by shift. The original iterator i will be transformed to i0 * ceil(n / nparts) + i1, where i0 and i1 are the iterators of the new outer and inner loops, respectively

Please note that the second mode will introduce an i0 * ceil(n / nparts) factor into the program, which cannot be recognized by polyhedral analysis, which may hinder some following schedules. If possible, plese use the first mode, and then reorder the inner and outer loops

Suppose the original loop is labeled "L", the split two loops can be selected by "$split.0{L}" (the outer loop) and "$split.1{L}" (the inner loop). If one of the resulting loop is proved to have only a single iteration, it will be removed

Parameters:	node (`str, ID or Stmt`) – The loop to be split factor (`int`) – Length of the inner loop. Set to -1 if using `nparts` nparts (`int`) – Length of the outer loop. Set to -1 if using `factor`

Exceptions:	`InvalidSchedule` – if the loop is not found

Returns:	`(Optional[ID], Optional[ID])` – (outer loop ID, inner loop ID), either ID can be None if the loop is proved to have only a single iteration

`freetensor.core.schedule.Schedule.swap(self, order)` ¶

Swap statements in the same block

To reorder nested loops, use reorder instead

Parameters:	order (`List[str (Selector string), ID, List[ID], Stmt, or List[Stmt]]`) – The statements. If one item of the `order` list contains multiple statements, the `order` list will be flattened

Exceptions:	`InvalidSchedule` – if the statements are not found or the dependences cannot be solved

`freetensor.core.schedule.Schedule.unroll(self, loop, immediate=False)` ¶

Unroll a loop

Parameters:

loop (str, ID or Stmt) – ID of the loop
immediate (bool) – If false (by default), postpone the unroll procedure to the backend compiler, which saves scheduling time. If true, unroll the loop immediately, which may help further simplifications based on the unrolled result. If your purpose is just to fill the instruction cache, set it to false. If you are unrolling a loop that computes array indices, set it to true

Exceptions:	`InvalidSchedule` – if the loop is not found or length of the loop is not a constant

`freetensor.core.schedule.Schedule.var_merge(self, vardef, dim)` ¶

Merge two dimensions of a variable

Parameters:	vardef (`str, ID or Stmt`) – ID of the VarDef statement of the specific variable dim (`int`) – Merge the `dim`-th and the `(dim + 1)`-th dimension

`freetensor.core.schedule.Schedule.var_reorder(self, vardef, order)` ¶

Reorder the dimensions of a variable

Parameters:	vardef (`str, ID or Stmt`) – ID of the VarDef statement of the specific variable order (`array like of str, ID or Stmt`) – Vector of integers. The new order of the dimensions

Exceptions:	`InvalidSchedule` – if the variable or the order is illegal

`freetensor.core.schedule.Schedule.var_split(self, vardef, dim, mode, factor=-1, nparts=-1)` ¶

Split a dimension of a variable into two

Parameters:

vardef (str, ID or Stmt) – ID of the VarDef statement of the specific variable
dim (int) – which dimension to be split
mode (VarSplitMode) – When the dimension to split is not divisible by factor or nparts, the resulting shape may become larger. In FixedSize mode, the actual buffer size will not be changed, and gurads will be added to prevent out-of-bound accesses. In RelaxedSize mode, the buffer size may increase. The RelaxedSize mode cannot be applied to I/O variables
factor (int) – Length of the inner (higher no.) dimension. Set to -1 if using nparts
nparts (int) – Length of the outer (lower no.) loop. Set to -1 if using factor

Exceptions:	`InvalidSchedule` – if the variable or the dimension is not found

`freetensor.core.schedule.Schedule.var_squeeze(self, vardef, dim)` ¶

Remove a singleton (1-lengthed) dimension from a variable

This is a utility schedule, which can be used together with var_split, var_merge and/or var_reorder to transform a variable to a desired shape.

Parameters:	vardef (`str, ID or Stmt`) – ID of the VarDef statement of the specific variable dim (`int`) – Remove the `dim`-th dimension

Exceptions:	`InvalidSchedule` – if the variable is not found or the dimension is illegal

`freetensor.core.schedule.Schedule.var_unsqueeze(self, vardef, dim)` ¶

Insert a singleton (1-lengthed) dimension to a variable

This is a utility schedule, which can be used together with var_split, var_merge and/or var_reorder to transform a variable to a desired shape.

Parameters:	vardef (`str, ID or Stmt`) – ID of the VarDef statement of the specific variable dim (`int`) – Insert a singleton dimension at the `dim`-th dimension

Exceptions:	`InvalidSchedule` – if the variable is not found or the dimension is illegal

`freetensor.core.schedule.Schedule.vectorize(self, loop)` ¶

Vectorize a loop

Please note that, as vectorization is different from architecture to achitecture, the scheduler may or may not postpone it to the backend compiler. The vectorization is a best-effort schedule

Parameters:	loop (`str, ID or Stmt`) – ID of the loop

Exceptions:	`InvalidSchedule` – if the ID or name is not found, or the dependence requirement is not met

`freetensor.core.schedule.schedule(ast=None, callback=None, backward_callback=None, jit_cache=<function cache at 0x147b6469ca60>, verbose=0)` ¶

Apply any schedule on an AST through a user callback

Parameters:

ast (Func or Stmt) – The AST to schedule. If not specified, a partial function will be returned that cna be used as a decorator
callback (Callable[[freetensor.core.schedule.Schedule], NoneType]) – Specify what schedule(s) to do in this callback
backward_callback (Callable[[freetensor.core.schedule.Schedule], NoneType]) – Specify what schedule(s) to do for the backward function, if ast is returned from AD with attach_backward=True. Defaults to be the same with callback
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
verbose (int) – 0 = print nothing. 1 = print the final AST. 2 = print an AST after each schedule

Returns:	`Func or JITTemplate` – Return a Func for an AST if there is no JIT parameters. Return a JITTemplate that generates a Func if there is at least one

`freetensor.core.staging` ¶

A staging framework to support the FreeTensor frontend.

`freetensor.core.staging.AllowShortcutScope` `dataclass` ¶

Allow return scope.

This is a context manager that allows return in statically deterministic control flow.

`freetensor.core.staging.BreakException (Exception)` ¶

Exception to be raised by StagingOverload.break_stmt.

Breaks from a for loop.

`freetensor.core.staging.ContinueException (Exception)` ¶

Exception to be raised by StagingOverload.continue_stmt.

Continues a for loop.

`freetensor.core.staging.ReplaceAnnotations (NodeTransformer)` ¶

`freetensor.core.staging.ReplaceAnnotations.generic_visit(self, node)` `inherited` ¶

Called if no explicit visitor function exists for a node.

`freetensor.core.staging.ReplaceAnnotations.visit(self, node)` `inherited` ¶

Visit a node.

`freetensor.core.staging.ReturnException (Exception)` ¶

Exception to be raised by StagingOverload.return_stmt.

Holds a return value that will be passed through to the function wrapper.

`freetensor.core.staging.StagedAssignable (ABC)` ¶

`freetensor.core.staging.StagedAssignable.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.staging.StagedAssignable.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.staging.StagedAssignable.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.staging.StagedAssignable.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.staging.StagedAssignable.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedIterable (ABC)` ¶

`freetensor.core.staging.StagedIterable.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.staging.StagedIterable.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.staging.StagedIterable.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.staging.StagedIterable.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.staging.StagedIterable.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedPredicate (ABC)` ¶

`freetensor.core.staging.StagedPredicate.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.staging.StagedPredicate.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.staging.StagedPredicate.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.staging.StagedPredicate.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.staging.StagedPredicate.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedTypeAnnotation` ¶

`freetensor.core.staging.StagedTypeAnnotation.class (ABCMeta)` `inherited` ¶

`freetensor.core.staging.StagedTypeAnnotation.class.base (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

freetensor.core.staging.StagedTypeAnnotation.__class__.__base__.__instancecheck__(cls, instance) special ¶

Override for isinstance(instance, cls).

freetensor.core.staging.StagedTypeAnnotation.__class__.__base__.__new__(mcls, name, bases, namespace, **kwargs) special staticmethod ¶

Create and return a new object. See help(type) for accurate signature.

freetensor.core.staging.StagedTypeAnnotation.__class__.__base__.__subclasscheck__(cls, subclass) special ¶

Override for issubclass(subclass, cls).

freetensor.core.staging.StagedTypeAnnotation.__class__.__base__.register(cls, subclass) ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedTypeAnnotation.class.register(cls, subclass)` `inherited` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedTypeAnnotationMeta (ABCMeta)` ¶

`freetensor.core.staging.StagedTypeAnnotationMeta.base (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.staging.StagedTypeAnnotationMeta.base.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.staging.StagedTypeAnnotationMeta.base.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.staging.StagedTypeAnnotationMeta.base.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.staging.StagedTypeAnnotationMeta.base.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedTypeAnnotationMeta.register(cls, subclass)` `inherited` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagedUnpackAssignable (ABC)` ¶

`freetensor.core.staging.StagedUnpackAssignable.class (type)` `inherited` ¶

Metaclass for defining Abstract Base Classes (ABCs).

Use this metaclass to create an ABC. An ABC can be subclassed directly, and then acts as a mix-in class. You can also register unrelated concrete classes (even built-in classes) and unrelated ABCs as 'virtual subclasses' -- these and their descendants will be considered subclasses of the registering ABC by the built-in issubclass() function, but the registering ABC won't show up in their MRO (Method Resolution Order) nor will method implementations defined by the registering ABC be callable (not even via super()).

`freetensor.core.staging.StagedUnpackAssignable.class.instancecheck(cls, instance)` `special` ¶

Override for isinstance(instance, cls).

`freetensor.core.staging.StagedUnpackAssignable.class.new(mcls, name, bases, namespace, **kwargs)` `special` `staticmethod` ¶

Create and return a new object. See help(type) for accurate signature.

`freetensor.core.staging.StagedUnpackAssignable.class.subclasscheck(cls, subclass)` `special` ¶

Override for issubclass(subclass, cls).

`freetensor.core.staging.StagedUnpackAssignable.class.register(cls, subclass)` ¶

Register a virtual subclass of an ABC.

Returns the subclass, to allow usage as a class decorator.

`freetensor.core.staging.StagingError (Exception)` ¶

Error occurred during staging function execution (i.e. IR tree generation).

`freetensor.core.staging.StagingOverload` ¶

`freetensor.core.staging.StagingOverload.allow_shortcut_scope(self, allow)` ¶

Opens a scope that allows shortcut control flows in a statically deterministic

context. Need to be closed by with statement.

`freetensor.core.staging.StagingOverload.assert_stmt(self, test)` ¶

Assert staging tool.

`freetensor.core.staging.StagingOverload.assign_stmt(self, name, value)` ¶

Customized assign wrapper.

If value is instance of StagedAssignable, it's regarded as a customized assign behavior and gets executed with the assigned target variable name. This wrapper is used for initializing a variable.

`freetensor.core.staging.StagingOverload.break_stmt(self)` ¶

Break staging tool. Only allow break in static control flow.

`freetensor.core.staging.StagingOverload.continue_stmt(self)` ¶

Continue staging tool. Only allow continue in static control flow.

`freetensor.core.staging.StagingOverload.custom_attr(self, obj, attr)` ¶

Customized attribute accessor.

The framework first looks for a Python native attribute. If not found, it looks for this overloaded custom attribute resolver.

The default implementation provides no custom attribute. Can be overridden by subclasses.

Parameters:	obj (`Any`) – Object to access attribute. attr (`str`) – Attribute name.

Returns:	`Any` – The attribute value.

`freetensor.core.staging.StagingOverload.foreach(self, names, iter, body)` ¶

Customized foreach wrapper.

If value is instance of StagedIterable, its regarded as a customized foreach behavior and used to generate code for the python for loop. Otherwise, we try to execute the loop as usual.

`freetensor.core.staging.StagingOverload.functiondef_wrapper(self, filename, func)` ¶

Function definition wrapper.

This wrapper performs extra initialization and cleanup for function definition.

`freetensor.core.staging.StagingOverload.if_then_else_expr(self, predicate, then_expr, else_expr)` ¶

If-then-else expression staging tool.

`freetensor.core.staging.StagingOverload.if_then_else_stmt(self, predicate, then_body, else_body=None)` ¶

If-then-else statement staging tool.

When predicate is deterministic in staging, only one branch is generated. Otherwise, a If node in IR is generated.

`freetensor.core.staging.StagingOverload.load_attr(self, obj, attr)` ¶

Load attribute staging tool. Allows customization of reading attributes.

`freetensor.core.staging.StagingOverload.mark_position(self, lineno)` ¶

Code position handler.

Defaults to a no-op. Can be overridden by subclasses.

Parameters:	lineno (`int`) – Line number of the next statement.

`freetensor.core.staging.StagingOverload.metadata(self, content)` ¶

Metadata handler.

A metadata line is a comment starting with #! and followed by a metadata, represented as a string parameter.

Defaults to a no-op. Can be overridden by subclasses.

Parameters:	content (`str`) – The metadata content.

`freetensor.core.staging.StagingOverload.return_stmt(self, value, funcname)` ¶

Return staging tool. Only allow return in static control flow.

`freetensor.core.staging.StagingOverload.unpack_assign_stmt(self, names, values)` ¶

Customized assign wrapper for one or more targets.

If values is instance of StagedUnpackAssignable, it's regarded as a customized assign behavior and gets executed with all the assigned targets' names. Otherwise, it calls assign_stmt with each sub-assignments.

Please note that names can be nested tuples like ("a", ("b", "c")).

Please also note that names can also be a single string like "a" even if values is a tuple. There is no unpacking in this case

`freetensor.core.staging.StagingOverload.while_stmt(self, fpred, body)` ¶

While statement staging tool.

`freetensor.core.staging.TransformError (Exception)` ¶

Error occurred during AST transforming from python function to staging function

that generates IR tree.

`freetensor.core.staging.Transformer (NodeTransformer)` `dataclass` ¶

Transformer(filename: 'str', base_lineno: 'int', curr_func: 'str' = None, nonlocals: 'List[List[str]]' = None)

`freetensor.core.staging.Transformer.generic_visit(self, node)` `inherited` ¶

Called if no explicit visitor function exists for a node.

`freetensor.core.staging.Transformer.visit(self, node)` ¶

Visit a node.

`freetensor.core.staging.Transformer.visit_AnnAssign(self, old_node)` ¶

Rule:

x: Ty -> freetensor__annotate__x = annotate_stmt('x', Ty) if freetensor__annotate__x: x = freetensor__annotate__x(x): pure annotation

`freetensor.core.staging.Transformer.visit_Assign(self, old_node)` ¶

Rule:

lhs = rhs -> lhs = unpack_assign_stmt('lhs', rhs) x.lhs = rhs -> x.lhs = unpack_assign_stmt('lhs', rhs) a, (b, c) = (x, (y, z)) -> a, (b, c) = unpack_assign_stmt(('a', ('b', 'c')), (x, (y, z))) a = b = c -> a = unpack_assign_stmt('a', c); b = unpack_assign_stmt('b', c)

If unpack_assign_stmt is not overloaded, assign_stmt will be called for each item

`freetensor.core.staging.Transformer.visit_Compare(self, old_node)` ¶

Expand multiple comparison into and expression.

`freetensor.core.staging.Transformer.visit_For(self, old_node)` ¶

Rule:

for x in iter:
    body

->

def for_body(x):
    body
foreach('x', iter, for_body)

`freetensor.core.staging.Transformer.visit_If(self, old_node)` ¶

Rule:

if pred:
    body
else:
    orelse

->

def then_body():
    body
def else_body():
    orelse
if_then_else_stmt(pred, then_body, else_body)

`freetensor.core.staging.Transformer.visit_IfExp(self, old_node)` ¶

Rule: body if test else orelse -> if_then_else_expr(test, body, orelse)

`freetensor.core.staging.Transformer.visit_While(self, old_node)` ¶

Rule:

while pred:
    body

->

def while_body():
    body
while_stmt(lambda: pred, while_body)

`freetensor.core.staging.call_helper(callee, *args, **kwargs)` ¶

Call helper that generates a python AST Call node with given callee (overload

member) and arguments AST node.

`freetensor.core.staging.function_helper(name, args, body, nonlocals)` ¶

Function helper that generates a python AST FunctionDef node with given name,

arguments name, and body.

`freetensor.core.stmt` ¶

Facility to build AST statements

Classes and functions in this module are internally used by transformer to construct ASTs. They are also used by some internal tests. API of these classes and functions are subject to changes. End users are encouraged to use transformer, instead of this module.

Classes and functions in this module are all in BigCamel naming style, to distinguish from expressions in expr.py

`freetensor.core.stmt.Assert` ¶

Scope used to create an Assert node

This scope is internally used by transformer and tests

E.g.:

with Assert(i > 0):
    ... # Assertion body

`freetensor.core.stmt.Else` ¶

Scope used to create an else branch of an If node

This scope is internally used by transformer and tests

E.g.:

with If(i > 0):
    ... # True branch
with Else():
    ... # Else branch

`freetensor.core.stmt.For` ¶

Scope used to create a For node

This scope is internally used by transformer and tests

E.g.:

with For('i', 0, n) as i:
    ... # Loop body

`freetensor.core.stmt.If` ¶

Scope used to create an If node

This scope is internally used by transformer and tests

E.g.:

with If(i > 0):
    ... # Branch body

`freetensor.core.stmt.Invoke` ¶

Inlined invocation of another AST

Invoke is used as a scope (with Invoke(...) as returned_vars), so that variables returned by the callee can be used in the socpe

Invoke can be used for invoking a gradient function, which has already been lowered as an AST. Please note that once a user function has been lowered as an AST, the dimensionalities of its tensors get fixed. Therefore, to invoke ordinary user functions, please use inline in transformer instead, which supports generic types

`freetensor.core.stmt.NamedScope` ¶

Scope used to create an StmtSeq node with an explicit labels

E.g.:

with NamedScope():
    ... # body

This scope is used for testing only. StmtSeq nodes can be deleted in many lowering passes

`freetensor.core.stmt.UserGradStaged` ¶

Internal staged implementation of UserGrad

`freetensor.core.stmt.Any()` ¶

Create an Any node (only for testing)

Any nodes matches any statement nodes in ast.match

`freetensor.core.stmt.Eval(expr)` ¶

Create an Eval node

This scope is internally used by transformer and tests

`freetensor.core.stmt.MarkLabel(label)` ¶

Mark the ID of the following statement

This scope is internally used by transformer and tests

`freetensor.core.stmt.MarkVersion(tape_name, var)` ¶

Create an MarkVersion node (only for custom gradient)

This node is only used for custom gradient. See UserGrad.

`freetensor.core.stmt.VarDef(*args, **kvs)` ¶

A factory function that creates a VarDef or a series of nested VarDefs

This scope is internally used by transformer and tests

`freetensor.core.transform` ¶

`freetensor.core.transform.inline(func=None, src=None, fallback=None, default_dynamic_range=True, verbose=False)` ¶

Enable a user function to be called by a transformed function at run time

Parameters:

func (Python function) – The user function
src (str (Optional)) – The source code of func. This parameter is only required if the source code cannot be get automatically, e.g., if func is generated from an exec
default_dynamic_range (bool) – If True, the built-in range is replaced with freetensor.dynamic_range. Defaults to True
verbose (bool) – True to print the generated Python code that is used for transforming

`freetensor.core.transform.transform(func=None, default_dynamic_range=True, bind={}, jit_cache=<function cache at 0x147b6469ca60>, target=None, verbose=0)` ¶

Transform a user function to an AST

Parameters:

func (Python function) – The user function to transform. If not specified, a partial function will be returend, which can be used as a decorator
default_dynamic_range (bool) – If True, the built-in range is replaced with freetensor.dynamic_range. Defaults to True
bind (Mapping[str, Any]) – Bind some parameters to specific values before transformations. Accpeting a parameter-name-to-value dict.
jit_cache (Callable[[Callable], Callable]) – Function decorator used to cache JIT instances
target (Target) – If not None, set config.default_target when transforming. This affects the default memory type use to create variables from Var, empty and etc.
verbose (int) – 0 = print nothing. 1 = print the resulting AST. 2 = 1 + print the generated Python code that is used for transforming

Returns:	`Func or JITTemplate` – Return a Func for an AST if there is no JIT parameters. Return a JITTemplate that generates a Func if there is at least one

`freetensor.core.utils` ¶

`freetensor.core.utils.as_decorator(f)` ¶

Enable a multi-parameter function f to be used as a decorator

Suppose g = as_decorator(f), enable the following usages:

@g
def h(...):
    ...

@g(a=a, b=b, c=c)
def h(...):
    ...

Formally, g will have the same parameters as f. f's first parameter should be the function it decorate, say h, and may have other parameters with default values. If h is set when called, g will return the decorated function, just as f does. If h is not set, g will return an f's partial function with all other parameters set, and the partial function can then be decorate another h again.

`freetensor.libop` `special` ¶

`freetensor.libop.assign` ¶

`freetensor.libop.assign.add_to(*_args, **_kvs)` ¶

(Broadcasted) add to a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.assign.assign(*_args, **_kvs)` ¶

(Broadcasted) assign to a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.assign.floordiv_to(*_args, **_kvs)` ¶

(Broadcasted) rounding-towards-negative-infinity integer division (following Python convention, but not C) from a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.assign.mod_to(*_args, **_kvs)` ¶

(Broadcasted) modulo (results are non-negative, following Python convention, but not C) from a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.assign.mul_to(*_args, **_kvs)` ¶

(Broadcasted) multiply to a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.assign.sub_to(*_args, **_kvs)` ¶

(Broadcasted) subtract from a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.assign.truediv_to(*_args, **_kvs)` ¶

(Broadcasted) floating-point division from a tensor two another tensor

Parameters:	y (`VarRef`) – The target tensor x (`VarRef`) – The source tensor

`freetensor.libop.concat` ¶

`freetensor.libop.concat.concat(inputs, axis=0)` ¶

Concatenate a list of tensors into a single tensor on an existing axis (out-of-place)

All input tensors must have the same shape, except for the dimension size of the axis to concatenate on.

All input tensors must have the same data type and memory type.

Parameters:	inputs (`Sequence[freetensor.core.expr.VarRef]`) – Tensors for concatenation axis (`int`) – Dimension number for concatenation. Negative axis means counting from the last dimension

Returns:	`VarRef` – Concatenation result

`freetensor.libop.concat.concat_(inputs, output, axis=0)` ¶

Concatenate a list of tensors into a single tensor on an existing axis (in-place)

All input tensors must have the same shape, except for the dimension size of the axis to concatenate on.

All input tensors must have the same data type and memory type.

Parameters:	inputs (`Sequence[freetensor.core.expr.VarRef]`) – Tensors for concatenation output (`VarRef`) – Concatenation result axis (`int`) – Dimension number for concatenation. Negative axis means counting from the last dimension

`freetensor.libop.concat.stack(inputs, axis=0)` ¶

Concatenate a list of tensors into a single tensor on a new axis (out-of-place)

All input tensors must have the same shape, data type and memory type.

Parameters:	inputs (`Sequence[freetensor.core.expr.VarRef]`) – Tensors for concatenation axis (`int`) – Dimension number for concatenation. Negative axis means counting from the last dimension

Returns:	`VarRef` – Concatenation result

`freetensor.libop.concat.stack_(inputs, output, axis=0)` ¶

Concatenate a list of tensors into a single tensor on a new axis (in-place)

All input tensors must have the same shape, data type and memory type.

Parameters:	inputs (`Sequence[freetensor.core.expr.VarRef]`) – Tensors for concatenation output (`VarRef`) – Concatenation result axis (`int`) – Dimension number for concatenation. Negative axis means counting from the last dimension

`freetensor.libop.constant` ¶

`freetensor.libop.constant.ones(shape, dtype, mtype=None)` ¶

Create a one-valued tensor

Parameters:	shape (`Sequence[Expr] or Var`) – Shape of the variable. A variable can be created using a literal shape, or another fixed-length VarRef as a shape dtype (`str or DataType`) – Data type of the variable mtype (`str or MemType (Optional)`) – Memory type of the variable. If omitted, the main memory type of the default Target in config will be used

Returns:	The one-valued tensor

`freetensor.libop.constant.ones_(y)` ¶

Fill ones to a tensor

Parameters:	y (`VarRef`) – The tensor to fill

`freetensor.libop.constant.zeros(shape, dtype, mtype=None)` ¶

Create a zero tensor

Parameters:	shape (`Sequence[Expr] or Var`) – Shape of the variable. A variable can be created using a literal shape, or another fixed-length VarRef as a shape dtype (`str or DataType`) – Data type of the variable mtype (`str or MemType (Optional)`) – Memory type of the variable. If omitted, the main memory type of the default Target in config will be used

Returns:	The zero tensor

`freetensor.libop.constant.zeros_(y)` ¶

Fill zeros to a tensor

Parameters:	y (`VarRef`) – The tensor to fill

`freetensor.libop.conv` ¶

`freetensor.libop.conv.conv(X, W, B=None, auto_pad='NOTSET', dilations=None, group=1, kernel_shape=None, pads=None, strides=None)` ¶

Convolution. The result is returned

Parameters follow ONNX convention. Currently only 2-D convolution is supported

`freetensor.libop.conv.conv_(X, W, B, Y, auto_pad='NOTSET', dilations=None, group=1, kernel_shape=None, pads=None, strides=None)` ¶

Convolution. The result is written to another tensor

Parameters follow ONNX convention. Currently only 2-D convolution is supported

`freetensor.libop.element_wise` ¶

`freetensor.libop.element_wise.abs(*_args, **_kvs)` ¶

Element-wise absolute value of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.abs_(*_args, **_kvs)` ¶

Element-wise absolute value of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.add(*_args, **_kvs)` ¶

(Broadcasted) element-wise addition of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.add_(*_args, **_kvs)` ¶

(Broadcasted) element-wise addition of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.binary_op(op, a, b)` ¶

(Broadcasted) any element-wise operation on two tensors and return the result

Parameters:	op (`Callable`) – The operation applied to each item a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.binary_op_(op, a, b, out)` ¶

(Broadcasted) any element-wise operation on two tensors. The result is written to another tensor

Parameters:	op (`Callable`) – The operation applied to each item a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.ceil(*_args, **_kvs)` ¶

Element-wise ceil of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.ceil_(*_args, **_kvs)` ¶

Element-wise ceil of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.ceildiv(*_args, **_kvs)` ¶

(Broadcasted) element-wise rounding-towards-positive-infinity integer division of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.ceildiv_(*_args, **_kvs)` ¶

(Broadcasted) element-wise rounding-towards-positive-infinity integer division of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.cos(*_args, **_kvs)` ¶

Element-wise cos of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.cos_(*_args, **_kvs)` ¶

Element-wise cos of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.eq(*_args, **_kvs)` ¶

(Broadcasted) element-wise equal of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.eq_(*_args, **_kvs)` ¶

(Broadcasted) element-wise equal of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.exp(*_args, **_kvs)` ¶

Element-wise natrual exponent of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.exp_(*_args, **_kvs)` ¶

Element-wise natrual exponent of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.floor(*_args, **_kvs)` ¶

Element-wise floor of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.floor_(*_args, **_kvs)` ¶

Element-wise floor of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.floordiv(*_args, **_kvs)` ¶

(Broadcasted) element-wise rounding-towards-negative-infinity integer division (following Python convention, but not C, recommended for performance) of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.floordiv_(*_args, **_kvs)` ¶

(Broadcasted) element-wise rounding-towards-negative-infinity integer division (following Python convention, but not C, recommended for performance) of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.ge(*_args, **_kvs)` ¶

(Broadcasted) element-wise greater-than-or-equal-to of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.ge_(*_args, **_kvs)` ¶

(Broadcasted) element-wise greater-than-or-equal-to of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.gt(*_args, **_kvs)` ¶

(Broadcasted) element-wise greater-than of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.gt_(*_args, **_kvs)` ¶

(Broadcasted) element-wise greater-than of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.l_and(*_args, **_kvs)` ¶

(Broadcasted) element-wise logical and of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.l_and_(*_args, **_kvs)` ¶

(Broadcasted) element-wise logical and of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.l_not(*_args, **_kvs)` ¶

Element-wise logical not of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.l_not_(*_args, **_kvs)` ¶

Element-wise logical not of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.l_or(*_args, **_kvs)` ¶

(Broadcasted) element-wise logical or of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.l_or_(*_args, **_kvs)` ¶

(Broadcasted) element-wise logical or of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.le(*_args, **_kvs)` ¶

(Broadcasted) element-wise less-than-or-equal-to of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.le_(*_args, **_kvs)` ¶

(Broadcasted) element-wise less-than-or-equal-to of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.ln(*_args, **_kvs)` ¶

Element-wise natrual logarithm of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.ln_(*_args, **_kvs)` ¶

Element-wise natrual logarithm of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.lt(*_args, **_kvs)` ¶

(Broadcasted) element-wise less-than of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.lt_(*_args, **_kvs)` ¶

(Broadcasted) element-wise less-than of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.max(*_args, **_kvs)` ¶

(Broadcasted) element-wise maximum of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.max_(*_args, **_kvs)` ¶

(Broadcasted) element-wise maximum of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.min(*_args, **_kvs)` ¶

(Broadcasted) element-wise minimum of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.min_(*_args, **_kvs)` ¶

(Broadcasted) element-wise minimum of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.mod(*_args, **_kvs)` ¶

(Broadcasted) element-wise modulo (results are non-negative, following Python convention, but not C, recommended for performance) of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.mod_(*_args, **_kvs)` ¶

(Broadcasted) element-wise modulo (results are non-negative, following Python convention, but not C, recommended for performance) of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.mul(*_args, **_kvs)` ¶

(Broadcasted) element-wise multiplication of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.mul_(*_args, **_kvs)` ¶

(Broadcasted) element-wise multiplication of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.ne(*_args, **_kvs)` ¶

(Broadcasted) element-wise non-equal of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.ne_(*_args, **_kvs)` ¶

(Broadcasted) element-wise non-equal of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.neg(*_args, **_kvs)` ¶

Element-wise negation of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.neg_(*_args, **_kvs)` ¶

Element-wise negation of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.relu(*_args, **_kvs)` ¶

Element-wise ReLU of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.relu_(*_args, **_kvs)` ¶

Element-wise ReLU of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.remainder(*_args, **_kvs)` ¶

(Broadcasted) element-wise remainder (results can be positive or negative, following C convention, but not Python, NOT recommended for performance) of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.remainder_(*_args, **_kvs)` ¶

(Broadcasted) element-wise remainder (results can be positive or negative, following C convention, but not Python, NOT recommended for performance) of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.round_towards_0_div(*_args, **_kvs)` ¶

(Broadcasted) element-wise rounding-towards-0 integer division (following C convention, but not Python, NOT recommended for performance) of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.round_towards_0_div_(*_args, **_kvs)` ¶

(Broadcasted) element-wise rounding-towards-0 integer division (following C convention, but not Python, NOT recommended for performance) of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.sigmoid(*_args, **_kvs)` ¶

Element-wise sigmoid of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.sigmoid_(*_args, **_kvs)` ¶

Element-wise sigmoid of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.sin(*_args, **_kvs)` ¶

Element-wise sin of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.sin_(*_args, **_kvs)` ¶

Element-wise sin of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.sqrt(*_args, **_kvs)` ¶

Element-wise square root of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.sqrt_(*_args, **_kvs)` ¶

Element-wise square root of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.square(*_args, **_kvs)` ¶

Element-wise square of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.square_(*_args, **_kvs)` ¶

Element-wise square of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.sub(*_args, **_kvs)` ¶

(Broadcasted) element-wise subtraction of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.sub_(*_args, **_kvs)` ¶

(Broadcasted) element-wise subtraction of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.tan(*_args, **_kvs)` ¶

Element-wise tan of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.tan_(*_args, **_kvs)` ¶

Element-wise tan of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.tanh(*_args, **_kvs)` ¶

Element-wise tanh of a tensor and return the result

Parameters:	x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.tanh_(*_args, **_kvs)` ¶

Element-wise tanh of a tensor. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.truediv(*_args, **_kvs)` ¶

(Broadcasted) element-wise floating-point division of two tensors and return the result

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.truediv_(*_args, **_kvs)` ¶

(Broadcasted) element-wise floating-point division of two tensors. The result is written to another tensor

Parameters:	a (`VarRef`) – Left-hand-side operand b (`VarRef`) – Right-hand-side operand out (`VarRef`) – The result tensor

`freetensor.libop.element_wise.unary_op(op, x)` ¶

Any element-wise operation on a tensor and return the result

Parameters:	op (`Callable`) – The operation applied to each item x (`VarRef`) – The input tensor

Returns:	`VarRef` – The result tensor

`freetensor.libop.element_wise.unary_op_(op, x, y)` ¶

Any element-wise operation on a tensor. The result is written to another tensor

Parameters:	op (`Callable`) – The operation applied to each item x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor

`freetensor.libop.logsumexp` ¶

`freetensor.libop.logsumexp.logsumexp(x, axis=-1, keepdims=True)` ¶

Compute ln sum_i exp(x_i), where i is along an axis. Return the result.

The computation is numerically stabilized.

Parameters:	x (`VarRef`) – The input tensor axis (`int`) – Axis that the reduction is performed along. Negative axis means counting from the last dimension keepdims (`bool`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.logsumexp.logsumexp_(x, y, axis=-1, keepdims=True)` ¶

Compute ln sum_i exp(x_i), where i is along an axis. Write to tensor y.

The computation is numerically stabilized.

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axis (`int`) – Axis that the reduction is performed along. Negative axis means counting from the last dimension keepdims (`bool`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.matmul` ¶

`freetensor.libop.matmul.einsum(fmt, *args)` ¶

Einstein summation. The result is returned

Parameters:	fmt (`str`) – The format string. E.g. `"ik,kj->ij"` represents a matrix multiplcation args (`Sequence[VarRef]`) – All inputs arguments. E.g. if `fmt` is `"ik,kj->ij"`, it iterates axis `i` and `k` of `args[0]`, axis `k` and `j` of `args[1]`, axis `i` and `j` of the returned value

Returns:	The result tensor

`freetensor.libop.matmul.einsum_(fmt, *args)` ¶

Einstein summation. The result is written to the last argument

Parameters:	fmt (`str`) – The format string. E.g. `"ik,kj->ij"` represents a matrix multiplcation args (`Sequence[VarRef]`) – All arguments including inputs and the output. E.g. if `fmt` is `"ik,kj->ij"`, it iterates axis `i` and `k` of `args[0]`, axis `k` and `j` of `args[1]`, axis `i` and `j` of `args[2]`

`freetensor.libop.matmul.gemm(A, B, C=None, has_bias=False, trans_A=False, trans_B=False, alpha=1.0, beta=1.0)` ¶

General matrix multiplcation following BLAS convention and return the result

It performs Y = alpha tr?(A) @ tr?(B) + C, where @ represents matrix multiplication, tr? represents an optional transposition

Parameters:

A (VarRef) – The left-hand-side operand of matrix multiplication
B (VarRef) – The right-hand-side operand of matrix multiplication
C (VarRef (Optional)) – The bias tensor
trans_A (bool) – If true, transpose A. Defaults to False
trans_B (bool) – If true, transpose B. Defaults to False
alpha (float) – Coefficient of tr?(A) @ tr?(B). Defaults to 1.0
beta (float) – Coefficient of C. Defaults to 1.0

Returns:	The resulting tensor

`freetensor.libop.matmul.gemm_(A, B, C, Y, trans_A=False, trans_B=False, alpha=1.0, beta=1.0)` ¶

General matrix multiplcation following BLAS convention. The result is written to an existing tensor

It performs Y = alpha tr?(A) @ tr?(B) + C, where @ represents matrix multiplication, tr? represents an optional transposition

Parameters:

A (VarRef) – The left-hand-side operand of matrix multiplication
B (VarRef) – The right-hand-side operand of matrix multiplication
C (VarRef (Optional)) – The bias tensor
Y (VarRef) – The resulting tensor
trans_A (bool) – If true, transpose A. Defaults to False
trans_B (bool) – If true, transpose B. Defaults to False
alpha (float) – Coefficient of tr?(A) @ tr?(B). Defaults to 1.0
beta (float) – Coefficient of C. Defaults to 1.0

`freetensor.libop.matmul.matmul(A, B)` ¶

Matrix multiplcation. The result is returned

Parameters:	A (`VarRef`) – The left-hand-side operand B (`VarRef`) – The right-hand-side operand

Returns:	The resulting tensor

`freetensor.libop.matmul.matmul_(A, B, Y)` ¶

Matrix multiplcation. The result is written to an existing tensor

Parameters:	A (`VarRef`) – The left-hand-side operand B (`VarRef`) – The right-hand-side operand C (`VarRef`) – The resulting tensor

`freetensor.libop.pooling` ¶

`freetensor.libop.pooling.global_avg_pool(X)` ¶

Global averaging pooling. The result is returned

Parameters follow ONNX convention. Currently only 2-D pooling is supported

`freetensor.libop.pooling.global_avg_pool_(X, Y)` ¶

Global averaging pooling. The result is written to another tensor

Parameters follow ONNX convention. Currently only 2-D pooling is supported

`freetensor.libop.pooling.max_pool(X, auto_pad='NOTSET', dilations=None, kernel_shape=None, pads=None, strides=None)` ¶

Maximum pooling. The result is returned

Parameters follow ONNX convention. Currently only 2-D pooling is supported

`freetensor.libop.pooling.max_pool_(X, Y, auto_pad='NOTSET', dilations=None, kernel_shape=None, pads=None, strides=None)` ¶

Maximum pooling. The result is written to another tensor

Parameters follow ONNX convention. Currently only 2-D pooling is supported

`freetensor.libop.reduction` ¶

`freetensor.libop.reduction.all(*_args, **_kvs)` ¶

Reduction of logical and of a tensor through one or more dimensions and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.reduction.all_(*_args, **_kvs)` ¶

Reduction of logical and of a tensor through one or more dimensions. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.reduction.any(*_args, **_kvs)` ¶

Reduction of logical or of a tensor through one or more dimensions and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.reduction.any_(*_args, **_kvs)` ¶

Reduction of logical or of a tensor through one or more dimensions. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.reduction.reduce_max(x, axes=None, keepdims=True)` ¶

Maximum of a tensor through one or more dimensions and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int]`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.reduction.reduce_max_(x, y, axes=None, keepdims=True)` ¶

Maximum of a tensor through one or more dimensions. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axes (`Sequence[int]`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.reduction.reduce_min(x, axes=None, keepdims=True)` ¶

Minimum of a tensor through one or more dimensions and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int]`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.reduction.reduce_min_(x, y, axes=None, keepdims=True)` ¶

Minimum of a tensor through one or more dimensions. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axes (`Sequence[int]`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.reduction.reduce_prod(*_args, **_kvs)` ¶

Product of a tensor through one or more dimensions and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.reduction.reduce_prod_(*_args, **_kvs)` ¶

Product of a tensor through one or more dimensions. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.reduction.reduce_sum(*_args, **_kvs)` ¶

Sum of a tensor through one or more dimensions and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

Returns:	`VarRef` – The result tensor

`freetensor.libop.reduction.reduce_sum_(*_args, **_kvs)` ¶

Sum of a tensor through one or more dimensions. The result is written to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axes (`Sequence[int] (Optional)`) – Which dimensions to reduce through. Defaults to None, standing for all dimensions, i.e., reduce the tensor to a scalar. Negative axis means counting form the last dimension keepdims (`bool (Optional)`) – Keep the reduced dimensions as singleton dimensions. Defaults to True

`freetensor.libop.reshape` ¶

`freetensor.libop.reshape.expand(a, expand_shape)` ¶

Broadcast a tensor to a given shape, following the broadcasting rules

Parameters:	a (`VarRef`) – The input tensor b (`Sequence of expressions`) – The broadcasted shape

Returns:	The broadcasted tensor

`freetensor.libop.reshape.expand_(a, out)` ¶

Broadcast a tensor to an existing tensor, following the broadcasting rules

Parameters:	a (`VarRef`) – The input tensor b (`VarRef`) – The broadcasted tensor

`freetensor.libop.reshape.flatten(x, axis=1)` ¶

Flatten a tensor to have two dimensions, and return the result

NOTE: This function follows the ONNX convension that reshapes to 2-D instead of 1-D.

Parameters:	x (`VarRef`) – The input tensor axis (`int`) – The result tensor will have 2 dimensions. All dimensions up to `axis` (inclusive) will be flattend to the first dimension. All dimensions after `axis` (exclusive) will be flatten to the second dimension. Negative axis means counting form the last dimension

Returns:	`VarRef` – The result tensor

`freetensor.libop.reshape.flatten_(x, y, axis=1)` ¶

Flatten a tensor to have two dimensions, and write to another tensor

NOTE: This function follows the ONNX convension that reshapes to 2-D instead of 1-D.

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axis (`int`) – The result tensor will have 2 dimensions. All dimensions up to `axis` (inclusive) will be flattend to the first dimension. All dimensions after `axis` (exclusive) will be flatten to the second dimension. Negative axis means counting form the last dimension

`freetensor.libop.reshape.flatten_onnx(x, axis=1)` ¶

Flatten a tensor to have two dimensions, and return the result

NOTE: This function follows the ONNX convension that reshapes to 2-D instead of 1-D.

Parameters:	x (`VarRef`) – The input tensor axis (`int`) – The result tensor will have 2 dimensions. All dimensions up to `axis` (inclusive) will be flattend to the first dimension. All dimensions after `axis` (exclusive) will be flatten to the second dimension. Negative axis means counting form the last dimension

Returns:	`VarRef` – The result tensor

`freetensor.libop.reshape.flatten_onnx_(x, y, axis=1)` ¶

Flatten a tensor to have two dimensions, and write to another tensor

NOTE: This function follows the ONNX convension that reshapes to 2-D instead of 1-D.

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axis (`int`) – The result tensor will have 2 dimensions. All dimensions up to `axis` (inclusive) will be flattend to the first dimension. All dimensions after `axis` (exclusive) will be flatten to the second dimension. Negative axis means counting form the last dimension

`freetensor.libop.reshape.flatten_pytorch(x, start_dim=0, end_dim=-1)` ¶

Flatten a tensor to have fewer dimensions, and return the result

NOTE: This function follows the PyTorch convension

Parameters:	x (`VarRef`) – The input tensor start_dim, end_dim (`int (Optional)`) – All dimensions ranging from `start_dim` and `end_dim` (inclusive) will be flattend to 1-D. Negative axis means counting form the last dimension

Returns:	`VarRef` – The result tensor

`freetensor.libop.reshape.flatten_pytorch_(x, y, start_dim=0, end_dim=-1)` ¶

Flatten a tensor to have fewer dimensions, and write to another tensor

NOTE: This function follows the PyTorch convension

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor start_dim, end_dim (`int (Optional)`) – All dimensions ranging from `start_dim` and `end_dim` (inclusive) will be flattend to 1-D. Negative axis means counting form the last dimension

`freetensor.libop.reshape.reshape(x, shape)` ¶

Reshape a tensor into a different shape with the same size

This operator will try to generate nested loops instead of looping over all elements in a plain loop, so schedules can be better applied. It guarantees to generates loops in the following cases:

Splitting a dimension. E.g. 4 to 2x2, and there will be a 2x2 loop nest.
Merging dimensions. E.g. 2x2 to 4, and there will be a 2x2 loop nest.
Each non-affecting dimension will be iterated by a unique loop. E.g. 3x5x7 to 5x3x7, and there will be a 15x7 loop nest, where the "7" dimension will be iterated by a unique loop.

Parameters:	x (`VarRef`) – The input tensor shape (`list of expression`) – The target shape

Returns:	`VarRef` – The result tensor

`freetensor.libop.reshape.reshape_(x, y)` ¶

Fill a tensor into another tensor with the same size but maybe different shape

This operator will try to generate nested loops instead of looping over all elements in a plain loop, so schedules can be better applied. It guarantees to generates loops in the following cases:

Splitting a dimension. E.g. 4 to 2x2, and there will be a 2x2 loop nest.
Merging dimensions. E.g. 2x2 to 4, and there will be a 2x2 loop nest.
Each non-affecting dimension will be iterated by a unique loop. E.g. 3x5x7 to 5x3x7, and there will be a 15x7 loop nest, where the "7" dimension will be iterated by a unique loop.

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor

`freetensor.libop.reshape.squeeze(x, axes)` ¶

Remove singleton dimensions from a tensor, and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int]`) – Dimension numbers of the singleton dimensions. Negative axis means counting from the last dimension

Returns:	`VarRef` – The resulting tensor

`freetensor.libop.reshape.squeeze_(x, y, axes)` ¶

Remove singleton dimensions from a tensor, and write the result to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The resulting tensor axes (`Sequence[int]`) – Dimension numbers of the singleton dimensions. Negative axis means counting from the last dimension

`freetensor.libop.reshape.unsqueeze(x, axes)` ¶

Insert singleton dimensions to a tensor, and return the result

Parameters:	x (`VarRef`) – The input tensor axes (`Sequence[int]`) – Dimension numbers of the new singleton dimensions. Negative axis means counting from the last dimension

Returns:	`VarRef` – The resulting tensor

`freetensor.libop.reshape.unsqueeze_(x, y, axes)` ¶

Insert singleton dimensions to a tensor, and write the result to another tensor

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The resulting tensor axes (`Sequence[int]`) – Dimension numbers of the new singleton dimensions. Negative axis means counting from the last dimension

`freetensor.libop.softmax` ¶

`freetensor.libop.softmax.softmax(x, axis=-1)` ¶

Softmax of tensor x along an axis and return the result

The computation is numerically stabilized.

Parameters:	x (`VarRef`) – The input tensor axis (`int (Optional)`) – Axis that the softmax is performed along. Negative axis means counting from the last dimension

Returns:	The result tensor

`freetensor.libop.softmax.softmax_(x, y, axis=-1)` ¶

Softmax of tensor x along an axis, and write to tensor y

The computation is numerically stabilized.

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The result tensor axis (`int`) – Axis that the softmax is performed along. Negative axis means counting from the last dimension

`freetensor.libop.transpose` ¶

`freetensor.libop.transpose.transpose(x, perm=None)` ¶

Transposition (out-of-place)

The perm[i]-th dimension of the input becomes the i-th dimension of the output

Parameters:	x (`VarRef`) – The input tensor perm (`Sequence[int]`) – Permutation of the dimensions. Negative values mean counting form the last dimension. By default reversing all dimensions

Returns:	`VarRef` – The output tensor

`freetensor.libop.transpose.transpose_(x, y, perm=None)` ¶

Transposition (in-place)

The perm[i]-th dimension of the input becomes the i-th dimension of the output

Parameters:	x (`VarRef`) – The input tensor y (`VarRef`) – The output tensor perm (`Sequence[int]`) – Permutation of the dimensions. Negative values mean counting form the last dimension. By default reversing all dimensions

Python API¶

freetensor.core special ¶

freetensor.core.autograd ¶

freetensor.core.autograd.grad(func=None, requires=None, provides=None, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, reset_provided_grad=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None) ¶

freetensor.core.autograd.grad_(func=None, requires=None, provides=None, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, reset_provided_grad=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None) ¶

freetensor.core.autograd.grad_body(stmt, requires, provides, tapes=<GradTapeMode.NoReuseOnly: 2>, reset_provided_grad=True, invert=False, user_grads=[]) ¶

freetensor.core.autograd.jacrev(func=None, inputs=None, output=None, flatten=False, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None) ¶

freetensor.core.autograd.jacrev_(func=None, inputs=None, output=None, flatten=False, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None) ¶

freetensor.core.codegen ¶

freetensor.core.codegen.NativeCode (EnableAttachBackward, NativeCode) ¶

freetensor.core.codegen.NativeCode.__contains__(self, item) special ¶

freetensor.core.codegen.codegen(ast=None, target=None, jit_cache=<function cache at 0x147b6469ca60>, verbose=False) ¶

freetensor.core.config ¶

freetensor.core.config.backend_compiler_cxx(*args, **kvs) ¶

freetensor.core.config.backend_compiler_nvcc(*args, **kvs) ¶

freetensor.core.config.backend_openmp(*args, **kvs) ¶

freetensor.core.config.debug_binary(*args, **kvs) ¶

freetensor.core.config.debug_cuda_with_um(*args, **kvs) ¶

freetensor.core.config.default_device(*args, **kvs) ¶

freetensor.core.config.default_target(*args, **kvs) ¶

freetensor.core.config.fast_math(*args, **kvs) ¶

freetensor.core.config.pretty_print(*args, **kvs) ¶

freetensor.core.config.print_all_id(*args, **kvs) ¶

freetensor.core.config.print_source_location(*args, **kvs) ¶

freetensor.core.config.runtime_dir(*args, **kvs) ¶

freetensor.core.config.set_backend_compiler_cxx(*args, **kvs) ¶

freetensor.core.config.set_backend_compiler_nvcc(*args, **kvs) ¶

freetensor.core.config.set_backend_openmp(*args, **kvs) ¶

freetensor.core.config.set_debug_binary(*args, **kvs) ¶

freetensor.core.config.set_debug_cuda_with_um(*args, **kvs) ¶

freetensor.core.config.set_default_device(*args, **kvs) ¶

freetensor.core.config.set_default_target(*args, **kvs) ¶

freetensor.core.config.set_fast_math(*args, **kvs) ¶

freetensor.core.config.set_pretty_print(*args, **kvs) ¶

freetensor.core.config.set_print_all_id(*args, **kvs) ¶

freetensor.core.config.set_print_source_location(*args, **kvs) ¶

freetensor.core.config.set_runtime_dir(*args, **kvs) ¶

freetensor.core.config.set_werror(*args, **kvs) ¶

freetensor.core.config.werror(*args, **kvs) ¶

freetensor.core.config.with_cuda(*args, **kvs) ¶

freetensor.core.config.with_mkl(*args, **kvs) ¶

freetensor.core.config.with_pytorch(*args, **kvs) ¶

freetensor.core.context ¶

freetensor.core.context.ContextStack ¶

freetensor.core.context.ContextStack.get_last_stmt_id(self) ¶

freetensor.core.context.ContextStack.push_append_stmt_callback(self, callback) ¶

freetensor.core.context.StmtRange ¶

freetensor.core.context.pop_ast(verbose=False) ¶

freetensor.core.context.pop_ast_and_user_grads(verbose=False) ¶

freetensor.core.driver ¶

freetensor.core.driver.Driver (EnableAttachBackward, Driver) ¶

freetensor.core.driver.Driver.__call__(self, *args, **kws) special ¶

freetensor.core.driver.Driver.__init__(self, native_code, device=None, host_device=None, cxx_flags=[], verbose=False) special ¶

freetensor.core.driver.Driver.collect_returns(self, always_return_pack=False) ¶

freetensor.core.driver.Driver.native_code(self) ¶

freetensor.core.driver.Driver.set_args(self, *args, **kws) ¶

freetensor.core.driver.array(data, dtype=None, dont_drop_borrow=False, moved=False) ¶

freetensor.core.driver.build_binary(code=None, device=None, host_device=None, jit_cache=<function cache at 0x147b6469ca60>, cxx_flags=[], verbose=False) ¶

freetensor.core.driver.move(data) ¶

freetensor.core.enable_attach_backward ¶

freetensor.core.enable_attach_backward.EnableAttachBackward ¶

freetensor.core.enable_attach_backward.EnableAttachBackward.__init__(self, *args, **kvs) special ¶

freetensor.core.expr ¶

freetensor.core.expr.AlreadyMadeReduceTo ¶

freetensor.core.expr.UndeclaredParam dataclass ¶

freetensor.core.expr.VarRef (FrontendVar) ¶

freetensor.core.expr.VarRef.as_reduce_to(self, reduce_op, metadata, value, atomic=False) ¶

freetensor.core.expr.VarRef.as_store(self, metadata, value) ¶

freetensor.core.expr.VarRef.select(self, idx, dim) ¶

freetensor.core.expr.VarRef.select_slice(self, begin, end, *, dim) ¶

freetensor.core.expr.VarRef.shape(self, dim=None) ¶

freetensor.core.expr.VarRefFromVarDef (VarRef) ¶

freetensor.core.expr.VarRefFromVarDef.as_reduce_to(self, reduce_op, metadata, value, atomic=False) inherited ¶

freetensor.core.expr.VarRefFromVarDef.as_store(self, metadata, value) inherited ¶

freetensor.core.expr.VarRefFromVarDef.select(self, idx, dim) inherited ¶

freetensor.core.expr.VarRefFromVarDef.select_slice(self, begin, end, *, dim) inherited ¶

freetensor.core.expr.VarRefFromVarDef.shape(self, dim=None) inherited ¶

freetensor.core.expr.VarVersionRef (VarRef) ¶

freetensor.core.expr.VarVersionRef.as_reduce_to(self, reduce_op, metadata, value, atomic=False) inherited ¶

freetensor.core.expr.VarVersionRef.as_store(self, metadata, value) inherited ¶

`freetensor.core` `special` ¶

`freetensor.core.autograd` ¶

`freetensor.core.autograd.grad(func=None, requires=None, provides=None, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, reset_provided_grad=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

`freetensor.core.autograd.grad_body(stmt, requires, provides, tapes=<GradTapeMode.NoReuseOnly: 2>, reset_provided_grad=True, invert=False, user_grads=[])` ¶

`freetensor.core.autograd.jacrev(func=None, inputs=None, output=None, flatten=False, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

`freetensor.core.autograd.jacrev_(func=None, inputs=None, output=None, flatten=False, tapes=<GradTapeMode.NoReuseOnly: 2>, tape_in_closure=True, invert=False, user_grads=None, attach_backward=False, jit_cache=<function cache at 0x147b6469ca60>, verbose=0, _override_func_params=None)` ¶

`freetensor.core.codegen` ¶

`freetensor.core.codegen.NativeCode (EnableAttachBackward, NativeCode)` ¶

`freetensor.core.codegen.NativeCode.contains(self, item)` `special` ¶

`freetensor.core.codegen.codegen(ast=None, target=None, jit_cache=<function cache at 0x147b6469ca60>, verbose=False)` ¶

`freetensor.core.config` ¶

`freetensor.core.config.backend_compiler_cxx(*args, **kvs)` ¶

`freetensor.core.config.backend_compiler_nvcc(*args, **kvs)` ¶

`freetensor.core.config.backend_openmp(*args, **kvs)` ¶

`freetensor.core.config.debug_binary(*args, **kvs)` ¶

`freetensor.core.config.debug_cuda_with_um(*args, **kvs)` ¶

`freetensor.core.config.default_device(*args, **kvs)` ¶

`freetensor.core.config.default_target(*args, **kvs)` ¶

`freetensor.core.config.fast_math(*args, **kvs)` ¶

`freetensor.core.config.pretty_print(*args, **kvs)` ¶

`freetensor.core.config.print_all_id(*args, **kvs)` ¶

`freetensor.core.config.print_source_location(*args, **kvs)` ¶

`freetensor.core.config.runtime_dir(*args, **kvs)` ¶

`freetensor.core.config.set_backend_compiler_cxx(*args, **kvs)` ¶

`freetensor.core.config.set_backend_compiler_nvcc(*args, **kvs)` ¶

`freetensor.core.config.set_backend_openmp(*args, **kvs)` ¶

`freetensor.core.config.set_debug_binary(*args, **kvs)` ¶

`freetensor.core.config.set_debug_cuda_with_um(*args, **kvs)` ¶

`freetensor.core.config.set_default_device(*args, **kvs)` ¶

`freetensor.core.config.set_default_target(*args, **kvs)` ¶

`freetensor.core.config.set_fast_math(*args, **kvs)` ¶

`freetensor.core.config.set_pretty_print(*args, **kvs)` ¶

`freetensor.core.config.set_print_all_id(*args, **kvs)` ¶

`freetensor.core.config.set_print_source_location(*args, **kvs)` ¶

`freetensor.core.config.set_runtime_dir(*args, **kvs)` ¶

`freetensor.core.config.set_werror(*args, **kvs)` ¶

`freetensor.core.config.werror(*args, **kvs)` ¶

`freetensor.core.config.with_cuda(*args, **kvs)` ¶

`freetensor.core.config.with_mkl(*args, **kvs)` ¶

`freetensor.core.config.with_pytorch(*args, **kvs)` ¶

`freetensor.core.context` ¶

`freetensor.core.context.ContextStack` ¶

`freetensor.core.context.ContextStack.get_last_stmt_id(self)` ¶

`freetensor.core.context.ContextStack.push_append_stmt_callback(self, callback)` ¶

`freetensor.core.context.StmtRange` ¶

`freetensor.core.context.pop_ast(verbose=False)` ¶

`freetensor.core.context.pop_ast_and_user_grads(verbose=False)` ¶

`freetensor.core.driver` ¶

`freetensor.core.driver.Driver (EnableAttachBackward, Driver)` ¶

`freetensor.core.driver.Driver.call(self, *args, **kws)` `special` ¶

`freetensor.core.driver.Driver.init(self, native_code, device=None, host_device=None, cxx_flags=[], verbose=False)` `special` ¶

`freetensor.core.driver.Driver.collect_returns(self, always_return_pack=False)` ¶

`freetensor.core.driver.Driver.native_code(self)` ¶

`freetensor.core.driver.Driver.set_args(self, *args, **kws)` ¶

`freetensor.core.driver.array(data, dtype=None, dont_drop_borrow=False, moved=False)` ¶

`freetensor.core.driver.build_binary(code=None, device=None, host_device=None, jit_cache=<function cache at 0x147b6469ca60>, cxx_flags=[], verbose=False)` ¶

`freetensor.core.driver.move(data)` ¶

`freetensor.core.enable_attach_backward` ¶

`freetensor.core.enable_attach_backward.EnableAttachBackward` ¶

`freetensor.core.enable_attach_backward.EnableAttachBackward.init(self, *args, **kvs)` `special` ¶

`freetensor.core.expr` ¶

`freetensor.core.expr.AlreadyMadeReduceTo` ¶

`freetensor.core.expr.UndeclaredParam` `dataclass` ¶

`freetensor.core.expr.VarRef (FrontendVar)` ¶

`freetensor.core.expr.VarRef.as_reduce_to(self, reduce_op, metadata, value, atomic=False)` ¶

`freetensor.core.expr.VarRef.as_store(self, metadata, value)` ¶

`freetensor.core.expr.VarRef.select(self, idx, dim)` ¶

`freetensor.core.expr.VarRef.select_slice(self, begin, end, *, dim)` ¶

`freetensor.core.expr.VarRef.shape(self, dim=None)` ¶

`freetensor.core.expr.VarRefFromVarDef (VarRef)` ¶

`freetensor.core.expr.VarRefFromVarDef.as_reduce_to(self, reduce_op, metadata, value, atomic=False)` `inherited` ¶

`freetensor.core.expr.VarRefFromVarDef.as_store(self, metadata, value)` `inherited` ¶

`freetensor.core.expr.VarRefFromVarDef.select(self, idx, dim)` `inherited` ¶

`freetensor.core.expr.VarRefFromVarDef.select_slice(self, begin, end, *, dim)` `inherited` ¶

`freetensor.core.expr.VarRefFromVarDef.shape(self, dim=None)` `inherited` ¶

`freetensor.core.expr.VarVersionRef (VarRef)` ¶

`freetensor.core.expr.VarVersionRef.as_reduce_to(self, reduce_op, metadata, value, atomic=False)` `inherited` ¶

`freetensor.core.expr.VarVersionRef.as_store(self, metadata, value)` `inherited` ¶

`freetensor.core.expr.VarVersionRef.select(self, idx, dim)` `inherited` ¶

`freetensor.core.expr.VarVersionRef.select_slice(self, begin, end, *, dim)` `inherited` ¶