Calc#
- class liesel.model.nodes.Calc(function, *inputs, _name='', _needs_seed=False, update_on_init=True, **kwinputs)[source]#
Bases:
Node
A
Node
subclass that calculates its value based on its inputs nodes.Calculator nodes are a central element of the Liesel graph building toolkit. They wrap arbitrary calculations in pure JAX functions.
By default, calculator nodes will appear in the node graph created by
viz.plot_nodes()
, but they will not appear in the model graph created byviz.plot_vars()
.You can use
new_calc()
if you want your calculation to be treated as a model variable and thus be shown inviz.plot_vars()
.
Tip
The wrapped function must be jit-compilable by JAX. This mainly means that it must be a pure function, i.e. it must not have any side effects and, given the same input, it must always return the same output. Some special consideration is also required for loops and conditionals.
Please consult the JAX docs for details.
- Parameters:
function (
Callable
[...
,Any
]) – The function to be wrapped. Must be jit-compilable by JAX.*inputs (
Any
) – Non-keyword inputs. Any inputs that are not already nodes orVar
will be converted toValue
nodes. The values of these inputs will be passed to the wrapped function in the same order they are entered here._name (
str
) – The name of the node. If you do not specify a name, a unique name will be automatically generated upon initialization of aModel
. (default:''
)_needs_seed (
bool
) – Whether the node needs a seed / PRNG key. (default:False
)update_on_init (
bool
) – IfTrue
, the calculator will try to evaluate its function upon initialization. (default:True
)**kwinputs (
Any
) – Keyword inputs. Any inputs that are not already nodes orVar`s will be converted to :class:
.Value` nodes. The values of these inputs will be passed to the wrapped function as keyword arguments.
See also
Var.new_calc
Initializes a weak variable that is a function of other variables.
Var
A variable in a statistical model, typically with a probability distribution.
Var.new_param
Initializes a strong variable that acts as a model parameter.
Var.new_obs
Initializes a strong variable that holds observed data.
Var.new_value
Initializes a strong variable without a distribution.
Value
A node representing some static data.
Dist
A node representing a
tensorflow_probability
Distribution
.
Examples
A simple calculator node, taking the exponential value of an input parameter.
>>> log_scale = lsl.Var.new_param(0.0, name="log_scale") >>> scale = lsl.Calc(jnp.exp, log_scale) >>> print(scale.value) 1.0
The value of the calculator node is updated when
Calc.update()
is called.>>> scale.update() Calc(name="") >>> print(scale.value) 1.0
You can also use your own functions as long as they are jit-compilable by JAX.
>>> def compute_variance(x): ... return jnp.exp(x)**2 >>> log_scale = lsl.Var.new_param(0.0, name="log_scale") >>> variance = lsl.Calc(compute_variance, log_scale).update() >>> print(variance.value) 1.0
Methods
add_inputs
(*inputs, **kwinputs)Adds non-keyword and keyword input nodes to the existing ones.
Returns all non-keyword and keyword input nodes as a unique tuple.
Returns all output nodes as a unique tuple.
Clears the state of the node.
Flags the node and its recursive outputs as outdated.
set_inputs
(*inputs, **kwinputs)Sets the non-keyword and keyword input nodes.
update
()Updates the value of the node.
Attributes
The wrapped function.
The groups that this node is a part of.
The non-keyword input nodes.
The keyword input nodes.
The model the node is part of.
The name of the node.
Whether the node needs a seed / PRNG key.
Whether the node is outdated.
The output nodes.
The state of the node.
The value of the node.
The variable the node is part of.
Whether the node should be monitored by an inference algorithm.