    def __call__(self, conv_filters: jnp.ndarray, conditioning: jnp.ndarray):
        """Applies FiLM conditioning to a convolutional feature map.

        Args:
            conv_filters: A tensor of shape [batch_size, height, width, channels].
            conditioning: A tensor of shape [batch_size, conditioning_size].

        Returns:
            A tensor of shape [batch_size, height, width, channels].
        """
        projected_cond_add = nn.Dense(
            features=conv_filters.shape[-1],
            kernel_init=nn.initializers.zeros,
            bias_init=nn.initializers.zeros,
        )(conditioning)
        projected_cond_mult = nn.Dense(
            features=conv_filters.shape[-1],
            kernel_init=nn.initializers.zeros,
            bias_init=nn.initializers.zeros,
        )(conditioning)

        projected_cond_add = projected_cond_add[..., None, None, :]
        projected_cond_mult = projected_cond_mult[..., None, None, :]

        return conv_filters * (1 + projected_cond_add) + projected_cond_mult
Applies FiLM conditioning to a convolutional feature map.

Args:
conv_filters: A tensor of shape [batch_size, height, width, channels].
conditioning: A tensor of shape [batch_size, conditioning_size].

Returns:
A tensor of shape [batch_size, height, width, channels].
serl_launcher/serl_launcher/vision/film_conditioning_layer.py