code: def batch_dense(inputs,
                  units,
                  activation=None,
                  kernel_initializer=None,
                  reuse=None,
                  name=None):
    
    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
docs: 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.