import tensorflow as tf

def baseline_loss(rewards, weights, average_across_timesteps=False, average_across_batch=True):
    """
    :param rewards: tensor of shape (batch_size, time_steps)
    :param weights: tensor of shape (batch_size, time_steps)
    """
    batch_size = tf.shape(rewards)[0]

    cost = rewards ** 2
    cost = tf.reduce_sum(cost * weights, axis=1)

    if average_across_timesteps:
        total_size = tf.reduce_sum(weights, axis=1)
        total_size += 1e-12  # just to avoid division by 0 for all-0 weights
        cost /= total_size

    cost = tf.reduce_sum(cost)

    if average_across_batch:
        cost /= tf.to_float(batch_size)

    return cost
