tf.Module

Base neural network module class.

Used in the notebooks

Used in the guide Used in the tutorials

A module is a named container for tf.Variables, other tf.Modules and functions which apply to user input. For example a dense layer in a neural network might be implemented as a tf.Module:

class Dense(tf.Module):
  def __init__(self, input_dim, output_size, name=None):
    super(Dense, self).__init__(name=name)
    self.w = tf.Variable(
      tf.random.normal([input_dim, output_size]), name='w')
    self.b = tf.Variable(tf.zeros([output_size]), name='b')
  def __call__(self, x):
    y = tf.matmul(x, self.w) + self.b
    return tf.nn.relu(y)

You can use the Dense layer as you would expect:

d = Dense(input_dim=3, output_size=2)