Calc

Contents

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 by viz.plot_vars().

  • You can use new_calc() if you want your calculation to be treated as a model variable and thus be shown in viz.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 or Var will be converted to Value 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 a Model. (default: '')

  • _needs_seed (bool) – Whether the node needs a seed / PRNG key. (default: False)

  • update_on_init (bool) – If True, the calculator will try to evaluate its function upon initialization. (default: True)

  • **kwinputs (Any) – Keyword inputs. Any inputs that are not already nodes or Var`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.

all_input_nodes()

Returns all non-keyword and keyword input nodes as a unique tuple.

all_output_nodes()

Returns all output nodes as a unique tuple.

clear_state()

Clears the state of the node.

flag_outdated()

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

function

The wrapped function.

groups

The groups that this node is a part of.

inputs

The non-keyword input nodes.

kwinputs

The keyword input nodes.

model

The model the node is part of.

name

The name of the node.

needs_seed

Whether the node needs a seed / PRNG key.

outdated

Whether the node is outdated.

outputs

The output nodes.

state

The state of the node.

value

The value of the node.

var

The variable the node is part of.

monitor

Whether the node should be monitored by an inference algorithm.