    def compute_output_tensor(self, input_tensor):
        """Returns the resulting tensor when all operations of this layer are applied to input_tensor.

        If input_tensor is 2D, this method returns a 3D tensor representing the output of each
        layer in the ensemble on the input_tensor. Otherwise, if the input_tensor is 3D, the output
        is also 3D, where output[i] = layer_ensemble[i](input[i]).

        Arguments:
            input_tensor: (tf.Tensor) The input to the layer.

        Returns: The output of the layer, as described above.
        """
        # Get raw layer outputs
        if len(input_tensor.shape) == 2:
            raw_output = tf.einsum("ij,ajk->aik", input_tensor, self.weights) + self.biases
        elif len(input_tensor.shape) == 3 and input_tensor.shape[0].value == self.ensemble_size:
            raw_output = tf.matmul(input_tensor, self.weights) + self.biases
        else:
            raise ValueError("Invalid input dimension.")

        # Apply activations if necessary
        return FC._activations[self.activation](raw_output)
Returns the resulting tensor when all operations of this layer are applied to input_tensor.

If input_tensor is 2D, this method returns a 3D tensor representing the output of each
layer in the ensemble on the input_tensor. Otherwise, if the input_tensor is 3D, the output
is also 3D, where output[i] = layer_ensemble[i](input[i]).

Arguments:
input_tensor: (tf.Tensor) The input to the layer.

Returns: The output of the layer, as described above.
mbpo/models/fc.py