<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 def dropout_with_broadcast_dims(x, keep_prob, broadcast_dims=None, **kwargs):
 """Like tf.nn.dropout but takes broadcast_dims instead of noise_shape. Instead of specifying noise_shape, this function takes broadcast_dims - a list of dimension numbers in which noise_shape should be 1. The random keep/drop tensor has dimensionality 1 along these dimensions. Args: x: a floating point tensor. keep_prob: A scalar Tensor with the same type as x. The probability that each element is kept. broadcast_dims: an optional list of integers the dimensions along which to broadcast the keep/drop flags. **kwargs: keyword arguments to tf.nn.dropout other than "noise_shape". Returns: Tensor of the same shape as x. """

  assert "noise_shape" not in kwargs
  if broadcast_dims:
    shape = tf.shape(x)
    ndims = len(x.get_shape())
    # Allow dimensions like "-1" as well.
    broadcast_dims = [dim + ndims if dim < 0 else dim for dim in broadcast_dims]
    kwargs["noise_shape"] = [
        1 if i in broadcast_dims else shape[i] for i in range(ndims)
    ]
  return tf.nn.dropout(x, keep_prob, **kwargs)