class_docstr: Wrapper to defer initialization of a `tf.Module` instance.
  
  `DeferredModule` is a general-purpose mechanism for creating objects that are
  'tape safe', meaning that computation occurs only when an instance
  method is called, not at construction. This ensures that method calls inside
  of a `tf.GradientTape` context will produce gradients to any underlying
  `tf.Variable`s.
  
  ### Examples
  
  TFP's built-in Distributions and Bijectors are tape-safe by contract, but
  this does not extend to cases where computation is required
  to construct an object's parameters prior to initialization.
  For example, suppose we want to construct a Gamma
  distribution with a given mean and variance. In a naive implementation,
  we would convert these to the Gamma's native `concentration` and
  `rate` parameters when the distribution is constructed. Any future method
  calls would produce gradients to `concentration` and `rate`, but not to the
  underlying mean and variance:
  
  ```python
  mean, variance = tf.Variable(3.2), tf.Variable(9.1)
  dist = tfd.Gamma(concentration=mean**2 / variance,
                   rate=mean / variance)
  
  with tf.GradientTape() as tape:
    lp = dist.log_prob(5.0)
  grads = tape.gradient(lp, [mean, variance])
  # ==> `grads` are `[None, None]` !! :-(
  ```
  
  To preserve the gradients, we can defer the parameter transformation using
  `DeferredModule`. The resulting object behaves just like a
  `tfd.Gamma` instance, however, instead of running the `Gamma` constructor just
  once, it internally applies the parameter transformation and constructs a
  new, temporary instance of `tfd.Gamma` on *every method invocation*.
  This ensures that all operations needed to compute a method's return value
  from any underlying variables are performed every time the method is invoked.
  A surrounding `GradientTape` context will therefore be able to trace the full
  computation.
  
  ```python
  def gamma_from_mean_and_variance(mean, variance, **kwargs):
    rate = mean / variance
    return tfd.Gamma(concentration=mean * rate, rate=rate, **kwargs)
  
  mean, variance = tf.Variable(3.2), tf.Variable(9.1)
  deferred_dist = tfp.experimental.util.DeferredModule(
    build_fn=gamma_from_mean_and_variance,
    mean=mean,  # May be passed by position or by name.
    variance=variance)
  
  with tf.GradientTape() as tape:
    lp = deferred_dist.log_prob(5.0)
  grads = tape.gradient(lp, [mean, variance])
  # ==> `grads` are defined!
  ```
  
  Note that we could have achieved a similar effect by using
  `tfp.util.DeferredTensor` to individually defer the `concentration` and `rate`
  parameters. However, this would have been significantly more verbose, and
  would not share any computation between the two parameter transformations.
  In general, `DeferredTensor` is often idiomatic for simple transformations of
  a single value, while `DeferredModule` may be preferred for transformations
  that operate on multiple values and/or contain multiple steps.
  
  ### Caveats
  
  Objects derived from a `DeferredModule` are no longer deferred, so
  they will not preserve gradients. For example, slicing into a deferred
  Distribution yields a new, concrete Distribution instance:
  
  ```python
  def normal_from_log_scale(scaled_loc, log_scale):
    return tfd.Normal(loc=5 * scaled_loc, scale=tf.exp(log_scale))
  
  dist = tfp.experimental.util.DeferredModule(
    build_fn=normal_from_log_scale,
    scaled_loc=tf.Variable([1., 2., 3.]),
    log_scale=tf.Variable([1., 1., 1.]))
  dist.batch_shape  # ==> [3]
  len(dist.trainable_variables)  # ==> 2
  
  slice = dist[:2]  # Instantiates a new, non-deferred Distribution.
  slice.batch_shape  # ==> [2]
  len(slice.trainable_variables)  # ==> 0 (!)
  
  # If needed, we could defer the slice with another layer of wrapping.
  deferred_slice = tfp.experimental.util.DeferredModule(
    build_fn=lambda d: d[:2],
    d=dist)
  len(deferred_slice.trainable_variables)  # ==> 2
  ```
docstr: Defers initialization of an object with transformed arguments.
  
  Args:
    build_fn: Python callable specifying a deferred transformation of the
      provided arguments. This must have signature
      `module = build_fn(*args, **kwargs)`. The return value `module` is an
      instance of `tf.Module`.
    *args: Optional positional arguments to `build_fn`.
    also_track: Optional instance or structure of instances of `tf.Variable`
      and/or `tf.Module`, containing any additional trainable variables that
      the `build_fn` may access beyond the given `args` and `kwargs`. This
      ensures that such variables will be correctly tracked in
      `self.trainable_variables`.
      Default value: `None`.
    **kwargs: Optional keyword arguments to `build_fn`.
code: def __init__(self, build_fn, *args, also_track=None, **kwargs):
      """Defers initialization of an object with transformed arguments.
  
      Args:
        build_fn: Python callable specifying a deferred transformation of the
          provided arguments. This must have signature
          `module = build_fn(*args, **kwargs)`. The return value `module` is an
          instance of `tf.Module`.
        *args: Optional positional arguments to `build_fn`.
        also_track: Optional instance or structure of instances of `tf.Variable`
          and/or `tf.Module`, containing any additional trainable variables that
          the `build_fn` may access beyond the given `args` and `kwargs`. This
          ensures that such variables will be correctly tracked in
          `self.trainable_variables`.
          Default value: `None`.
        **kwargs: Optional keyword arguments to `build_fn`.
      """
      self._build_fn = build_fn
      self._param_args = args
      self._param_kwargs = kwargs
      self._deferred_module_also_track = also_track
  
      # In order for DeferredModule to work as a tf.Module, we need to ensure that
      # attrs used by tf.Module are handled directly, rather than being forwarded
      # to the inner class.
      self._module_attrs = set(dir(tf.Module()))
      super(DeferredModule, self).__init__()
  
file_path: probability/tensorflow_probability/python/experimental/util/deferred_module.py
imports: import tensorflow.compat.v2
  from tensorflow_probability.python.experimental.util import special_methods
label_desc: Use image node_id 1 to create a new DeferredModule object from inherited base classes: tf, special_methods with example: obj = DeferredModule(build_fn)