View source on GitHub

A class to store nested objects in a collection of numpy arrays.

If a data_spec of {'foo': ArraySpec(shape=(4,), dtype=np.uint8), 'bar': ArraySpec(shape=(3, 7), dtype=np.float32)} were used, then this would create two arrays, one for the 'foo' key and one for the 'bar' key. The .get and .set methods would return/take Python dictionaries, but break down the component arrays before storing them.

data_spec An ArraySpec or a list/tuple/nest of ArraySpecs describing a single item that can be stored in this table.
capacity The maximum number of items that can be stored in the buffer.

ValueError If data_spec is not an instance or nest of ArraySpecs.

name Returns the name of this module as passed or determined in the ctor.

name_scope Returns a tf.name_scope instance for this class.
submodules Sequence of all sub-modules.

Submodules are modules which are properties of this module, or found as properties of modules which are properties of this module (and so on).

a = tf.Module()
b = tf.Module()
c = tf.Module()
a.b = b
b.c = c
list(a.submodules) == [b, c]
list(b.submodules) == [c]
list(c.submodules) == []

trainable_variables Sequence of trainable variables owned by this module and its submodules.

variables Sequence of variables owned by this module and its submodules.



View source

Get value stored at idx.


View source

Set table_idx to value.


Decorator to automatically enter the module name scope.

class MyModule(tf.Module):
  def __call__(self, x):
    if not hasattr(self, 'w'):
      self.w = tf.Variable(tf.random.normal([x.shape[1], 3]))
    return tf.matmul(x, self.w)

Using the above module would produce tf.Variables and tf.Tensors whose names included the module name:

mod = MyModule()
mod(tf.ones([1, 2]))
<tf.Tensor: shape=(1, 3), dtype=float32, numpy=..., dtype=float32)>
<tf.Variable 'my_module/Variable:0' shape=(2, 3) dtype=float32,
numpy=..., dtype=float32)>

method The method to wrap.

The original method wrapped such that it enters the module's name scope.