Home
Overview
Fundamentals
Glossary
Cheatsheet
API Reference
Quick Recipes
Deep Dive
Common Pitfalls and Gotchas
Persistent Storage
Operators to create and call functors.
kd.functor.FunctorFactory(*args, **kwargs)`functor_factory` argument protocol for `trace_as_fn`.
Implements:
(py_types.FunctionType, return_type_as: arolla.QValue) -> DataItemkd.functor.SWITCH_DEFAULTMarks the default case in a kd.switch() construct.kd.functor.allow_arbitrary_unused_inputs(fn_def: DataItem) -> DataItemReturns a functor that allows unused inputs but otherwise behaves the same.
This is done by adding a `**__extra_inputs__` argument to the signature if
there is no existing variadic keyword argument there. If there is a variadic
keyword argument, this function will return the original functor.
This means that if the functor already accepts arbitrary inputs but fails
on unknown inputs further down the line (for example, when calling another
functor), this method will not fix it. In particular, this method has no
effect on the return values of kd.py_fn or kd.bind. It does however work
on the output of kd.trace_py_fn.
Args:
fn_def: The input functor.
Returns:
The input functor if it already has a variadic keyword argument, or its copy
but with an additional `**__extra_inputs__` variadic keyword argument if
there is no existing variadic keyword argument.kd.functor.bind(fn_def: DataItem, /, *args: Any, return_type_as: Any = <class 'koladata.types.data_slice.DataSlice'>, **kwargs: Any) -> DataItemAliases:
Returns a Koda functor that partially binds a function.
This function is intended to work the same as functools.partial in Python.
The bound positional arguments are prepended to the arguments provided at call
time. For keyword arguments, for every "k=something" argument that you pass
to this function, whenever the resulting functor is called, if the user did
not provide "k=something_else" at call time, we will add "k=something".
Moreover, if the user provides a value for a positional-or-keyword argument
positionally, and it was previously bound using bind(..., k=v), an exception
will occur.
You can pass expressions with their own inputs as values in `args` and
`kwargs`. Those inputs will become inputs of the resulting functor, will be
used to compute those expressions, _and_ they will also be passed to the
underying functor.
Example:
f = kd.bind(kd.fn(I.x + I.y), x=0)
kd.call(f, y=1) # 1
g = kd.bind(kd.fn(S + 7), 5)
kd.call(g) # 12
Args:
fn_def: A Koda functor.
*args: Positional arguments to bind. The values may be Koda expressions or
DataItems.
return_type_as: The return type of the functor is expected to be the same as
the type of this value. This needs to be specified if the functor does not
return a DataSlice. kd.types.DataSlice and kd.types.DataBag can also be
passed here.
**kwargs: Keyword arguments to bind. The values in this map may be Koda
expressions or DataItems. When they are expressions, they must evaluate to
a DataSlice/DataItem or a primitive that will be automatically wrapped
into a DataItem. This function creates auxiliary variables with names
starting with '_aux_fn', so it is not recommended to pass variables with
such names.
Returns:
A new Koda functor with some parameters bound.kd.functor.call(fn, *args, return_type_as=None, **kwargs)Aliases:
Calls a functor.
See the docstring of `kd.fn` on how to create a functor.
Example:
kd.call(kd.fn(I.x + I.y), x=2, y=3)
# returns kd.item(5)
kd.lazy.call(I.fn, x=2, y=3).eval(fn=kd.fn(I.x + I.y))
# returns kd.item(5)
Args:
fn: The functor to be called, typically created via kd.fn().
*args: The positional arguments to pass to the call. Scalars will be
auto-boxed to DataItems.
return_type_as: The return type of the call is expected to be the same as
the return type of this expression. In most cases, this will be a literal
of the corresponding type. This needs to be specified if the functor does
not return a DataSlice. kd.types.DataSlice, kd.types.DataBag and
kd.types.JaggedShape can also be passed here.
**kwargs: The keyword arguments to pass to the call. Scalars will be
auto-boxed to DataItems.
Returns:
The result of the call.kd.functor.call_and_update_namedtuple(fn, *args, namedtuple_to_update, **kwargs)Calls a functor which returns a namedtuple and applies it as an update.
This operator exists to avoid the need to specify return_type_as for the inner
call (since the returned namedtuple may have a subset of fields of the
original namedtuple, potentially in a different order).
Example:
kd.functor.call_and_update_namedtuple(
kd.fn(lambda x: kd.namedtuple(x=x * 2)),
x=2,
namedtuple_to_update=kd.namedtuple(x=1, y=2))
# returns kd.namedtuple(x=4, y=2)
Args:
fn: The functor to be called, typically created via kd.fn().
*args: The positional arguments to pass to the call. Scalars will be
auto-boxed to DataItems.
namedtuple_to_update: The namedtuple to be updated with the result of the
call. The returned namedtuple must have a subset (possibly empty or full)
of fields of this namedtuple, with the same types.
**kwargs: The keyword arguments to pass to the call. Scalars will be
auto-boxed to DataItems.
Returns:
The updated namedtuple.kd.functor.call_fn_normally_when_parallel(fn, *args, return_type_as=None, **kwargs)Special call that will invoke the functor normally in parallel mode.
Normally, nested functor calls are also parallelized in parallel mode. This
operator can be used to disable the nested parallelization for a specific
call. Instead, the parallel evaluation will first wait for all inputs of
the call to be ready, and then call the functor normally on them. Note that
the functor must accept and return non-parallel types (DataSlices, DataBags)
and not futures. The functor may return an iterable, which will be converted
to a stream in parallel mode, or another non-parallel value, which will
be converted to a future in parallel mode. The functor must not accept
iterables as inputs, which will result in an error.
Outside of the parallel mode, this operator behaves exactly like
`functor.call`.
Args:
fn: The functor to be called, typically created via kd.fn().
*args: The positional arguments to pass to the call.
return_type_as: The return type of the call is expected to be the same as
the return type of this expression. In most cases, this will be a literal
of the corresponding type. This needs to be specified if the functor does
not return a DataSlice. kd.types.DataSlice, kd.types.DataBag and
kd.types.JaggedShape can also be passed here.
**kwargs: The keyword arguments to pass to the call.
Returns:
The result of the call.kd.functor.call_fn_returning_stream_when_parallel(fn, *args, return_type_as=ITERABLE[DATA_SLICE]{sequence(value_qtype=DATA_SLICE)}, **kwargs)Special call that will be transformed to expect fn to return a stream.
It should be used only if functor is provided externally in production
enviroment. Prefer `functor.call` for functors fully implemented in Python.
Args:
fn: function to be called. Should return Iterable in interactive mode and
Stream in parallel mode.
*args: positional args to pass to the function.
return_type_as: The return type of the call is expected to be the same as
the return type of this expression. In most cases, this will be a literal
of the corresponding type. This needs to be specified if the functor does
not return a Iterable[DataSlice].
**kwargs: The keyword arguments to pass to the call.kd.functor.expr_fn(returns: Any, *, signature: DataSlice | None = None, auto_variables: bool = False, **variables: Any) -> DataItemCreates a functor.
Args:
returns: What should calling a functor return. Will typically be an Expr to
be evaluated, but can also be a DataItem in which case calling will just
return this DataItem, or a primitive that will be wrapped as a DataItem.
When this is an Expr, it either must evaluate to a DataSlice/DataItem, or
the return_type_as= argument should be specified at kd.call time.
signature: The signature of the functor. Will be used to map from args/
kwargs passed at calling time to I.smth inputs of the expressions. When
None, the default signature will be created based on the inputs from the
expressions involved.
auto_variables: When true, we create additional variables automatically
based on the provided expressions for 'returns' and user-provided
variables. All non-scalar-primitive DataSlice literals become their own
variables, and all named subexpressions become their own variables. This
helps readability and manipulation of the resulting functor.
**variables: The variables of the functor. Each variable can either be an
expression to be evaluated, or a DataItem, or a primitive that will be
wrapped as a DataItem. The result of evaluating the variable can be
accessed as V.smth in other expressions.
Returns:
A DataItem representing the functor.kd.functor.flat_map_chain(iterable, fn, value_type_as=None)Alias for kd.iterables.flat_map_chain
kd.functor.flat_map_interleaved(iterable, fn, value_type_as=None)Alias for kd.iterables.flat_map_interleaved
kd.functor.fn(f: Any, *, use_tracing: bool = True, **kwargs: Any) -> DataItemAliases:
Returns a Koda functor representing `f`.
This is the most generic version of the functools builder functions.
It accepts all functools supported function types including python functions,
Koda Expr.
Args:
f: Python function, Koda Expr, Expr packed into a DataItem, or a Koda
functor (the latter will be just returned unchanged).
use_tracing: Whether tracing should be used for Python functions.
**kwargs: Either variables or defaults to pass to the function. See the
documentation of `expr_fn` and `py_fn` for more details.
Returns:
A Koda functor representing `f`.kd.functor.for_(iterable, body_fn, *, finalize_fn=unspecified, condition_fn=unspecified, returns=unspecified, yields=unspecified, yields_interleaved=unspecified, **initial_state)Alias for kd.iterables.for_
kd.functor.fstr_fn(returns: str, **kwargs) -> DataItemReturns a Koda functor from format string.
Format-string must be created via Python f-string syntax. It must contain at
least one formatted expression.
kwargs are used to assign values to the functor variables and can be used in
the formatted expression using V. syntax.
Each formatted expression must have custom format specification,
e.g. `{I.x:s}` or `{V.y:.2f}`.
Examples:
kd.call(fstr_fn(f'{I.x:s} {I.y:s}'), x=1, y=2) # kd.slice('1 2')
kd.call(fstr_fn(f'{V.x:s} {I.y:s}', x=1), y=2) # kd.slice('1 2')
kd.call(fstr_fn(f'{(I.x + I.y):s}'), x=1, y=2) # kd.slice('3')
kd.call(fstr_fn('abc')) # error - no substitutions
kd.call(fstr_fn('{I.x}'), x=1) # error - format should be f-string
Args:
returns: A format string.
**kwargs: variable assignments.kd.functor.get_inspect_signature(fn_def: DataItem) -> SignatureRetrieves the signature of the given functor as an inspect.Signature.
This function retrieves the Koda Functor signature attached to the given
functor and converts it to a Python `inspect.Signature` object.
Args:
fn_def: The functor to retrieve the signature for.
Returns:
The Python signature derived from the functor's Koda signature.kd.functor.get_signature(fn_def: DataItem) -> DataItemRetrieves the signature attached to the given functor.
Args:
fn_def: The functor to retrieve the signature for, or a slice thereof.
Returns:
The signature(s) attached to the functor(s).kd.functor.has_fn(x)Aliases:
Returns `present` for each item in `x` that is a functor.
Note that this is a pointwise operator. See `kd.functor.is_fn` for the
corresponding scalar version.
Args:
x: DataSlice to check.kd.functor.if_(cond, yes_fn, no_fn, *args, return_type_as=None, **kwargs)Aliases:
Calls either `yes_fn` or `no_fn` depending on the value of `cond`.
This a short-circuit sibling of kd.cond which executes only one of the two
functors, and therefore requires that `cond` is a MASK scalar.
Example:
x = kd.item(5)
kd.if_(x > 3, lambda x: x + 1, lambda x: x - 1, x)
# returns 6, note that both lambdas were traced into functors.
Args:
cond: The condition to branch on. Must be a MASK scalar.
yes_fn: The functor to be called if `cond` is present.
no_fn: The functor to be called if `cond` is missing.
*args: The positional argument(s) to pass to the functor.
return_type_as: The return type of the call is expected to be the same as
the return type of this expression. In most cases, this will be a literal
of the corresponding type. This needs to be specified if the functor does
not return a DataSlice. kd.types.DataSlice, kd.types.DataBag and
kd.types.JaggedShape can also be passed here.
**kwargs: The keyword argument(s) to pass to the functor.
Returns:
The result of the call of either `yes_fn` or `no_fn`.kd.functor.is_fn(obj: Any) -> DataItemAliases:
Checks if `obj` represents a functor.
Args:
obj: The value to check.
Returns:
kd.present if `obj` is a DataSlice representing a functor, kd.missing
otherwise (for example if obj has wrong type).kd.functor.map(fn, *args, include_missing=False, **kwargs)Aliases:
Aligns fn and args/kwargs and calls corresponding fn on corresponding arg.
If certain items of `fn` are missing, the corresponding items of the result
will also be missing.
If certain items of `args`/`kwargs` are missing then:
- when include_missing=False (the default), the corresponding item of the
result will be missing.
- when include_missing=True, we are still calling the functor on those missing
args/kwargs.
`fn`, `args`, `kwargs` will all be broadcast to the common shape. The return
values of the functors will be converted to a common schema, or an exception
will be raised if the schemas are not compatible. In that case, you can add
the appropriate cast inside the functor.
Example:
fn1 = kdf.fn(lambda x, y: x + y)
fn2 = kdf.fn(lambda x, y: x - y)
fn = kd.slice([fn1, fn2])
x = kd.slice([[1, None, 3], [4, 5, 6]])
y = kd.slice(1)
kd.map(kd.slice([fn1, fn2]), x=x, y=y)
# returns kd.slice([[2, None, 4], [3, 4, 5]])
kd.map(kd.slice([fn1, None]), x=x, y=y)
# returns kd.slice([[2, None, 4], [None, None, None]])
Args:
fn: DataSlice containing the functor(s) to evaluate. All functors must
return a DataItem.
*args: The positional argument(s) to pass to the functors.
include_missing: Whether to call the functors on missing items of
args/kwargs.
**kwargs: The keyword argument(s) to pass to the functors.
Returns:
The evaluation result.kd.functor.map_py_fn(f: Union[Callable[..., Any], PyObject], *, schema: Any = None, max_threads: Any = 1, ndim: Any = 0, include_missing: Any = None, **defaults: Any) -> DataItemReturns a Koda functor wrapping a python function for kd.map_py.
See kd.map_py for detailed APIs, and kd.py_fn for details about function
wrapping. schema, max_threads and ndims cannot be Koda Expr or Koda functor.
Args:
f: Python function.
schema: The schema to use for resulting DataSlice.
max_threads: maximum number of threads to use.
ndim: Dimensionality of items to pass to `f`.
include_missing: Specifies whether `f` applies to all items (`=True`) or
only to items present in all `args` and `kwargs` (`=False`, valid only
when `ndim=0`); defaults to `False` when `ndim=0`.
**defaults: Keyword defaults to pass to the function. The values in this map
may be kde expressions, format strings, or 0-dim DataSlices. See the
docstring for py_fn for more details.kd.functor.map_reduce_update(fn, *args, include_missing=False, **kwargs)Aliases:
Aligns fn and args/kwargs and calls corresponding fn on corresponding arg.
Similar to `kd.map`, but expects each functor to return a DataBag instead of
a DataItem. All the returned DataBags are merged into a single DataBag.
Example:
fn = kdf.fn(lambda x: kd.attrs(x, y=x.a + 1))
x = kd.slice([kd.new(a=1), kd.new(a=2)])
update_bag = kd.map_reduce_update(fn, x=x)
# returns DataBag equal to kd.attrs(x, y=kd.slice([2, 3]))
Args:
fn: DataSlice containing the functor(s) to evaluate. All functors must
return a DataBag.
*args: The positional argument(s) to pass to the functors.
include_missing: Whether to call the functors on missing items of
args/kwargs.
**kwargs: The keyword argument(s) to pass to the functors.
Returns:
A DataBag containing the merged results of all functor evaluations.kd.functor.py_fn(f: Callable[..., Any], *, return_type_as: Any = <class 'koladata.types.data_slice.DataSlice'>, **defaults: Any) -> DataItemAliases:
Returns a Koda functor wrapping a python function.
This is the most flexible way to wrap a python function for large, complex
code that doesn't require serialization. Note that functions wrapped with
py_fn are not serializable. See register_py_fn for an alternative that is
serializable.
Note that unlike the functors created by kd.functor.expr_fn from an Expr, this
functor will have exactly the same signature as the original function. In
particular, if the original function does not accept variadic keyword
arguments and an unknown argument is passed when calling the functor, an
exception will occur.
Args:
f: Python function. It is required that this function returns a
DataSlice/DataItem or a primitive that will be automatically wrapped into
a DataItem.
return_type_as: The return type of the function is expected to be the same
as the type of this value. This needs to be specified if the function does
not return a DataSlice/DataItem or a primitive that would be auto-boxed
into a DataItem. kd.types.DataSlice, kd.types.DataBag and
kd.types.JaggedShape can also be passed here.
**defaults: Keyword defaults to bind to the function. The values in this map
may be Koda expressions or DataItems (see docstring for kd.bind for more
details). Defaults can be overridden through kd.call arguments. **defaults
and inputs to kd.call will be combined and passed through to the function.
If a parameter that is not passed does not have a default value defined by
the function then an exception will occur.
Returns:
A DataItem representing the functor.kd.functor.reduce(fn, items, initial_value)Alias for kd.iterables.reduce
kd.functor.register_py_fn(f: Callable[..., Any], *, return_type_as: Any = <class 'koladata.types.data_slice.DataSlice'>, unsafe_override: bool = False, **defaults: Any) -> DataItemAliases:
Returns a Koda functor wrapping a function registered as an operator.
This is the recommended way to wrap a non-traceable python function into a
functor.
`f` will be wrapped into an operator and registered with the name taken from
the `__module__` and `__qualname__` attributes. This requires `f` to be named,
and to not be a locally defined function. Furthermore, attempts to call
`register_py_fn` on the same function will fail unless unsafe_override is True
(not recommended).
The resulting functor can be serialized and loaded from a different process
(unlike `py_fn`). In order for the serialized functor to be deserialized, an
equivalent call to `register_py_fn` is required to have been made in the
process that performs the deserialization. In practice, this is often
implemented through `f` being a top-level function that is traced at library
import time.
Note that unlike the functors created by kd.functor.expr_fn from an Expr, this
functor will have exactly the same signature as the original function. In
particular, if the original function does not accept variadic keyword
arguments and an unknown argument is passed when calling the functor, an
exception will occur.
Args:
f: Python function. It is required that this function returns a
DataSlice/DataItem or a primitive that will be automatically wrapped into
a DataItem. The function must also have `__qualname__` and `__module__`
attributes set.
return_type_as: The return type of the function is expected to be the same
as the type of this value. This needs to be specified if the function does
not return a DataSlice/DataItem or a primitive that would be auto-boxed
into a DataItem. kd.types.DataSlice, kd.types.DataBag and
kd.types.JaggedShape can also be passed here.
unsafe_override: Whether to override an existing operator. Not recommended
unless you know what you are doing.
**defaults: Keyword defaults to bind to the function. The values in this map
may be Koda expressions or DataItems (see docstring for kd.bind for more
details). Defaults can be overridden through kd.call arguments. **defaults
and inputs to kd.call will be combined and passed through to the function.
If a parameter that is not passed does not have a default value defined by
the function then an exception will occur.
Returns:
A DataItem representing the functor.kd.functor.sub(functor: Any, replacements: dict[Any, Any], *, ignore_signature_checks: bool = False) -> AnyReplaces variables on the functor graph by their value or itemids.
This function traverses the entire tree of functor and its subfunctors and
replaces every occurrence of the keys in `replacements` with the corresponding
value, returning a new, modified functor. All appearances of a given value as
a variable in all sub-functors will be replaced.
Example 1:
f = kd.fn(lambda: kd.with_name(1, 'x') + kd.with_name(2, 'y'))
f1 = kd.functor.sub(f, {1: 100})
f() # Returns 3
f1() # Returns 102
Example 2:
@kd.trace_as_fn()
def mean(values):
return kd.math.agg_mean(values)
@kd.trace_as_fn()
def harmonic_mean(values):
return kd.math.agg_mean(values)
@kd.fn
def my_program():
...
# Replaces all uses of arithmetic mean with harmonic mean in my_program and
# its subfunctors.
new_program = kd.functor.sub(my_program, {mean: harmonic_mean})
If the value being replaced is not sufficiently unique (e.g. the number 1 in
example 1) this may end up doing unintended replacements that are hard to
debug. Consider this other example:
@kd.trace_as_fn()
def g(x):
return x + kd.with_name(1, 'p') # p has value 1
@kd.fn
def f(x):
return g(x) * kd.with_name(1, 'q') # q also has value 1
f(1) # returns (1 + 1) * 1 = 2
f = substitute.sub(f, {f.g.p: 10})
f(1) # returns (1 + 10) * 10 = 110, not 11.
Both f.q and f.g.p are named variables and have the same value, so both are
replaced. Therefore, for safety, if you know exactly the variable you want to
replace, it is safer to use with_attrs on the functor that needs the change.
For example:
f.with_attrs(foo=bar) # simple and preferred: updates the outermost functor.
f.updated(kd.attrs(f.child, qux=baz)) # updates a child functor, also good.
However, if the change you need is in a deeply nested functor, or it is a
variable that shows up in many places, this gets harder to write correctly. In
this case, kd.functor.sub can help, provided that the original value is
sufficiently unique.
Beware when replacing DataSlices. Functors are represented as scalar objects,
so all slices are imploded into lists during tracing, and exploded when read.
Therefore, we disallow replacement values that are not scalar DataItems. If
you need to replace a named DataSlice, pass a list as a replacement.
Args:
functor: the functor to be traversed.
replacements: a mapping from original values to their replacements.
ignore_signature_checks: If True, bypasses signature validation for
subfunctor replacements.
Returns:
A copy of the input functor with variables replaced.
Raises:
ValueError if a functor is being replaced with a functor that doesn't have a
compatible signature and ignore_signature_checks is True (the default). See
`kd.functor.assert_signatures_compatible` for details.kd.functor.sub_by_name(functor: Any, replacements: dict[str, Any] | list[tuple[str, Any]], *, ignore_signature_checks: bool = False) -> AnyReplaces variables on the functor graph, matching by their names.
Traverses the entire graph of a non-recursive functor, including subfunctors,
replacing variables on each functor with the provided replacements. Matching
is done by name, and can match both variables like kd.V.x and subfunctors.
Only variables are replaced. No inputs or arbitrary operations within the
functor are performed. All values passed as replacement must be eager values.
When replacing functors, we require that both the original and the replacement
have a compatible signature, so that the program structure remains valid.
Also, there must not be two different variables with the same name in the
functor subtree. In this case, we raise a ValueError, because both would be
pointing to the same new value, likely causing hard to debug issues.
Examples:
1. Replacing vars.
f = kd.fn(lambda y: kd.with_name(1, 'x') + y)
f_new = kd.functor.sub_by_name(f, {'x': 5})
f_new(1) # Returns 6
2. Replacing subfunctors.
@kd.trace_as_fn
def double(x):
return x * 2
@kd.trace_as_fn()
def halve(x):
return x / 2
f = kd.fn(lambda x: double(x) + 1)
f_new = kd.functor.sub_by_name(f, {'double': kd.fn(halve)})
f(10) # Returns 21
f_new(10) # Returns 6.0 (with schema kd.FLOAT32)
3. Variable name collision
g = kd.fn(lambda y: kd.with_name(1, 'x') + y)
h = kd.fn(lambda y: kd.with_name(1000, 'x') * y)
f = kd.fn(lambda y: g(y) + h(y))
kd.functor.sub_by_name(f, {'x': 5}) # ValueError: two different `x` found.
Args:
functor: The root functor to modify.
replacements: A dictionary or list of pairs mapping names to their new
values/functors.
ignore_signature_checks: If True, bypasses signature validation for
subfunctor replacements.
Returns:
The new functor with recursively updated dependencies.kd.functor.switch(key, cases, *args, return_type_as=None, **kwargs)Aliases:
Calls a functor selected from `cases` based on `key`.
Raises an error if `kd.SWITCH_DEFAULT` case is not provided and `key` is
missing in `cases`.
Example:
kd.switch(
'a',
{
'a': lambda x: x + 1,
'b': lambda x: x * 2
kd.SWITCH_DEFAULT: lambda **unused: 57,
},
x=10)
Args:
key: A DataSlice key to lookup in `cases`.
cases: A scalar (Python or Koda) dict of functors.
*args: Positional arguments to pass to the selected functor.
return_type_as: The return type of the call is expected to be the same as
the return type of this expression.
**kwargs: Keyword arguments to pass to the selected functor.
Returns:
The result of calling the selected functor.kd.functor.trace_as_fn(*, name: str | None = None, return_type_as: Any = None, functor_factory: FunctorFactory | None = None)Aliases:
A decorator to customize the tracing behavior for a particular function.
A function with this decorator is converted to an internally-stored functor.
In traced expressions that call the function, that functor is invoked as a
sub-functor via by 'kde.call', rather than the function being re-traced.
Additionally, the functor passed to 'kde.call' is assigned a name, so that
when auto_variables=True is used (which is the default in kd.trace_py_fn),
the functor for the decorated function will become an attribute of the
functor for the outer function being traced.
The result of 'kde.call' is also assigned a name with a '_result' suffix, so
that it also becomes an separate variable in the outer function being traced.
This is useful for debugging.
This can be used to avoid excessive re-tracing and recompilation of shared
python functions, to quickly add structure to the functor produced by tracing
for complex computations, or to conveniently embed a py_fn into a traced
expression.
When using kd.parallel.call_multithreaded, using this decorator on
sub-functors can improve parallelism, since all sub-functor calls
are treated as separate tasks to be parallelized there.
This decorator is intended to be applied to standalone functions.
When applying it to a lambda, consider specifying an explicit name, otherwise
it will be called '<lambda>' or '<lambda>_0' etc, which is not very useful.
When applying it to a class method, it is likely to fail in tracing mode
because it will try to auto-box the class instance into an expr, which is
likely not supported.
When executing the resulting function in eager mode, we will evaluate the
underlying function directly instead of evaluating the functor, to have
nicer stack traces in case of an exception. However, we will still apply
the boxing rules on the returned value (for example, convert Python primitives
to DataItems) to better emulate what will happen in tracing mode.kd.functor.trace_py_fn(f: Callable[..., Any], *, auto_variables: bool = True, **defaults: Any) -> DataItemAliases:
Returns a Koda functor created by tracing a given Python function.
When 'f' has variadic positional (*args) or variadic keyword
(**kwargs) arguments, their name must start with 'unused', and they
must actually be unused inside 'f'.
'f' must not use Python control flow operations such as if or for.
Args:
f: Python function.
auto_variables: When true, we create additional variables automatically
based on the traced expression. All DataSlice literals become their own
variables, and all named subexpressions become their own variables. This
helps readability and manipulation of the resulting functor. Note that
this defaults to True here, while it defaults to False in
kd.functor.expr_fn.
**defaults: Keyword defaults to bind to the function. The values in this map
may be Koda expressions or DataItems (see docstring for kd.bind for more
details). Defaults can be overridden through kd.call arguments. **defaults
and inputs to kd.call will be combined and passed through to the function.
If a parameter that is not passed does not have a default value defined by
the function then an exception will occur.
Returns:
A DataItem representing the functor.kd.functor.visit_functors(functor: DataItem, callback_fn: Callable[[DataItem, dict[str, Any]], Any]) -> DataItemRecursively visits functors in the functor graph, bottom-up.
Traverses the graph of nested functors. For each functor, `callback_fn`
is called with:
* the functor itself, and
* a dict mapping attribute names that hold subfunctors to their
already-visited (transformed) results.
To visit / transform functor variables (without recursing into subfunctors),
use kd.functor.visit_variables.
Example:
fn = kd.fn(kd.V.g(kd.I.x, g=kd.fn(lambda x: x + 1)))
def transform(fn, subfunctors):
# subfunctors == {'g': <visited g>}
return fn.with_attrs(**subfunctors)
visitor.visit_functors(fn, transform)
Args:
functor: Root functor to traverse.
callback_fn: Called on each functor with its visited subfunctors.
Returns:
Visit result.kd.functor.visit_variables(functor: DataItem, callback_fn: Callable[[str, DataItem, dict[str, Any]], Any]) -> AnyVisits variables of a single functor in topological order.
This function does **not** recurse into subfunctors. It iterates over the
variables of `functor` in evaluation order (topological order) and calls
`callback_fn` on each variable.
`callback_fn` receives:
* the variable name (str),
* the variable value (DataItem - kd.EXPR which can be accessed via
kd.expr.unpack_expr, while all other DataItems represent literals
(even subfunctors)),
* a dict mapping variable names to the results of all previously visited
variables.
Args:
functor: The functor whose variables to visit.
callback_fn: Called on each variable.
Returns:
A visit result for 'returns' variable.kd.functor.while_(condition_fn, body_fn, *, returns=unspecified, yields=unspecified, yields_interleaved=unspecified, **initial_state)Alias for kd.iterables.while_