kd.nn.interms_property

kauldron.modules.interms_property()

interms property that makes storing intermediates more convenient.

Usage:

class MyModule(nn.Module):
  interms = kd.nn.interms_property()  # No typing annotation here !

  @nn.compact
  def __call__(self, x):
    h = nn.Dense(12)(x)

    # append to intermediates.
    self.interms["hidden"] = h

    # Access with `interms.path.to.module.hidden[0]`  (no `__call__`)
    # Setting interms like above is equivalent to using:
    # self.sow("intermediates", "hidden", h)
    ...

    # access intermediates
    h = self.interms["hidden"])

    # The above is equivalent to:
    # self.get_variable("intermediates", "hidden")[-1]


    # access intermediates by path
    res = nn.Dense(128, name="submod")(h)

    # res can also be accessed by absolute path:
    self.interms.get_by_path("path.to.my_module.submod.__call__[0]")
    # by relative path:
    self.interms.get_by_path(".submod.__call__[0]")
    # or via the parent
    self.interms.get_by_path("..my_module.submod.__call__[0])

    return res

The interms property can only be used within .init / .apply.

Returns:

The interms property