Multiply a batch of input matrices by a batch of parameter matrices.

  Each input matrix is multiplied by the corresponding parameter matrix.

  This is useful in a mixture-of-experts where the batch represents different
  experts with different inputs.

  Args:
    inputs: a Tensor with shape [batch, length, input_units]
    units: an integer
    activation: an optional activation function to apply to the output
    kernel_initializer: an optional initializer
    reuse: whether to reuse the varaible scope
    name: an optional string

  Returns:
    a Tensor with shape [batch, length, units]

  Raises:
    ValueError: if the "batch" or "input_units" dimensions of inputs are not
      statically known.

def batch_dense(inputs,
                units,
                activation=None,
                kernel_initializer=None,
                reuse=None,
                name=None):
  """Multiply a batch of input matrices by a batch of parameter matrices.

  Each input matrix is multiplied by the corresponding parameter matrix.

  This is useful in a mixture-of-experts where the batch represents different
  experts with different inputs.

  Args:
    inputs: a Tensor with shape [batch, length, input_units]
    units: an integer
    activation: an optional activation function to apply to the output
    kernel_initializer: an optional initializer
    reuse: whether to reuse the varaible scope
    name: an optional string

  Returns:
    a Tensor with shape [batch, length, units]

  Raises:
    ValueError: if the "batch" or "input_units" dimensions of inputs are not
      statically known.
  """
  inputs_shape = shape_list(inputs)
  if len(inputs_shape) != 3:
    raise ValueError("inputs must have 3 dimensions")
  batch = inputs_shape[0]
  input_units = inputs_shape[2]
  if not isinstance(batch, int) or not isinstance(input_units, int):
    raise ValueError("inputs must have static dimensions 0 and 2")
  with tf.variable_scope(
      name,
      default_name="batch_dense",
      values=[inputs],
      reuse=reuse,
      dtype=inputs.dtype):
    if kernel_initializer is None:
      kernel_initializer = tf.random_normal_initializer(
          stddev=input_units**-0.5)
    w = tf.get_variable(
        "w", [batch, input_units, units],
        initializer=kernel_initializer,
        dtype=inputs.dtype)
    y = tf.matmul(inputs, w)
    if activation is not None:
      y = activation(y)
    return y