code: def dense_relu_dense(x,
                       hidden_channels,
                       dropout=0.0,
                       dropout_broadcast_dims=None,
                       master_dtype=tf.float32,
                       slice_dtype=tf.float32, name=None):
    
    with tf.variable_scope(name, default_name="dense_relu_dense"):
      io_channels = x.shape.dims[-1]
      h = dense(x, hidden_channels,
                use_bias=False, activation=mtf.relu,
                master_dtype=master_dtype, slice_dtype=slice_dtype, name="wi")
      if dropout != 0.0:
        h = mtf.dropout(h, 1.0 - dropout,
                        noise_shape=h.shape - dropout_broadcast_dims)
      return dense(h, io_channels, use_bias=False, activation=None,
                   master_dtype=master_dtype, slice_dtype=slice_dtype,
                   name="wo")
docs: Hidden layer with ReLU activation followed by linear projection.
  
    The output has the same number of channels as the input.
  
    Args:
      x: a mtf.Tensor
      hidden_channels: a mtf.Dimension - channels in the hidden layer
      dropout: an optional float
      dropout_broadcast_dims: an optional list of mtf.Dimension
      master_dtype: a tf.dtype
      slice_dtype: a tf.dtype
      name: an optional string
  
    Returns:
      a mtf.Tensor with the same shape as x.