class_docstr: Parallel linear.
  
  This is equivalent to n separate linears applied in parallel to n inputs. It
  takes an input of shape [num_linears, batch_size, input_size] and returns an
  output of shape [num_linears, batch_size, output_size].
  
  It uses a single batched matmul which is more efficient than stacking separate
  snt.Linear layers. This is implemented using `num_linear`s first to avoid the
  need for transposes in order to make it efficient when stacking these.
docstr: Constructs parameters used by this module.
code: def _initialize(self, inputs: tf.Tensor):
      """Constructs parameters used by this module."""
      utils.assert_rank(inputs, 3)
  
      self.input_size = inputs.shape[2]
      if self.input_size is None:  # Can happen inside an @tf.function.
          raise ValueError(
              "Input size must be specified at module build time."
          )
      num_linears = inputs.shape[0]
      if num_linears is None:  # Can happen inside an @tf.function.
          raise ValueError(
              "The number of linears must be specified at module build time."
          )
  
      if self.w_init is None:
          # See https://arxiv.org/abs/1502.03167v3.
          stddev = 1.0 / math.sqrt(self.input_size)
          self.w_init = initializers.TruncatedNormal(stddev=stddev)
  
      self.w = tf.Variable(
          self.w_init(
              [num_linears, self.input_size, self.output_size],
              inputs.dtype,
          ),
          name="w",
      )
  
      if self.with_bias:
          self.b = tf.Variable(
              self.b_init(
                  [num_linears, 1, self.output_size], inputs.dtype
              ),
              name="b",
          )
  
file_path: sonnet/sonnet/src/parallel_linear.py
imports: import math
  from typing import Optional
  from sonnet.src import base
  from sonnet.src import initializers
  from sonnet.src import once
  from sonnet.src import utils
  import tensorflow
label_desc: Use image node_id 2 for calling the ParallelLinears obj's underlying member method code with example usage: obj._initialize(inputs) without return types