# Path: telluride_decoding/brain_model.py
# def pearson_correlation(x, y):
#   def positive_fcn():
#   def negative_fcn():
# def pearson_correlation_first(x, y):
# def pearson_correlation_second(x, y):
#   def call(self, x, y):
#   def __init__(self, tensorboard_dir=None, **kwargs):
#   def tensorboard_dir(self):
#   def fit(self, x=None, y=None, **kwargs):
#   def evaluate(self, x=None, y=None, epoch_count=1, **kwargs):
#   def add_metadata(self, flags, dataset=None):
#   def add_tensorboard_summary(self, name, data, subdir='train', step=0):
#   def __init__(self, input_dataset, regularization_lambda=0.0,
#                tensorboard_dir=None, **kwargs):
#   def call(self, input_dataset):
#   def compile(self, optimizer=tf.keras.optimizers.RMSprop,
#               loss='mse',
#               metrics=pearson_correlation_first,
#               learning_rate=1e-3, **kwargs):
#   def fit(self, input_dataset, **kwargs):
#   def weight_matrices(self):
# def calculate_linear_regressor_parameters_from_dataset(dataset, lamb=0.1,
#                                                        use_offset=True,
#                                                        use_ridge=True):
#   def __init__(self, input_dataset, num_hidden_list=None, **kwargs):
#   def call(self, input_dataset):
#   def compile(self, optimizer=tf.keras.optimizers.RMSprop,
#               loss='mse',
#               metrics=(pearson_correlation_first, 'mse'),
#               learning_rate=1e-3, **kwargs):
#   def fit(self, input_dataset, **kwargs):
#   def __init__(self, input_dataset, num_hidden_list=None, **kwargs):
#   def compile(self, optimizer=tf.keras.optimizers.Adam,
#               loss=tf.keras.losses.BinaryCrossentropy(),
#               metrics='accuracy', learning_rate=1e-3, **kwargs):
#   def call(self, input_dataset):
#   def fit(self, input_dataset, **kwargs):
# class PearsonCorrelationLoss(tf.keras.losses.Loss):
# class BrainModel(tf.keras.models.Model):
# class BrainModelLinearRegression(BrainModel):
# class BrainModelDNN(BrainModel):
# class BrainModelClassifier(BrainModel):


from absl import logging
from telluride_decoding import brain_model
import numpy as np
import tensorflow.compat.v2 as tf
                                shape=(self.input1_dim, real_output_dims),
                                initializer=self._rot1_init,
                                trainable=False)
    self.rot2 = self.add_weight(name='rot2',
                                shape=(self.input2_dim, real_output_dims),
                                initializer=self._rot2_init,
                                trainable=False)
    logging.info('BrainCcaLayer.build mean1=%s, mean2=%s, rot1=%s, rot2=%s',
                 self.mean1.shape, self.mean2.shape,
                 self.rot1.shape, self.rot2.shape)

  def call(self, inputs):
    input1 = inputs[0]
    input2 = inputs[1]

    input1 = tf.reshape(input1, (-1, self.input1_dim))
    input2 = tf.reshape(input2, (-1, self.input2_dim))

    rotated_input1 = tf.keras.backend.dot(input1 - self.mean1, self.rot1)
    rotated_input2 = tf.keras.backend.dot(input2 - self.mean2, self.rot2)

    return tf.concat((rotated_input1, rotated_input2),
                     axis=1, name='ConcatenateCcaResults')

  def get_config(self):
    """Defines the parameters needed when re-creating this layer."""
    return {'requested_cca_dims': self.output_dims,
           }



# Path: telluride_decoding/brain_model.py
# def pearson_correlation(x, y):
#   def positive_fcn():
#   def negative_fcn():
# def pearson_correlation_first(x, y):
# def pearson_correlation_second(x, y):
#   def call(self, x, y):
#   def __init__(self, tensorboard_dir=None, **kwargs):
#   def tensorboard_dir(self):
#   def fit(self, x=None, y=None, **kwargs):
#   def evaluate(self, x=None, y=None, epoch_count=1, **kwargs):
#   def add_metadata(self, flags, dataset=None):
#   def add_tensorboard_summary(self, name, data, subdir='train', step=0):
#   def __init__(self, input_dataset, regularization_lambda=0.0,
#                tensorboard_dir=None, **kwargs):
#   def call(self, input_dataset):
#   def compile(self, optimizer=tf.keras.optimizers.RMSprop,
#               loss='mse',
#               metrics=pearson_correlation_first,
#               learning_rate=1e-3, **kwargs):
#   def fit(self, input_dataset, **kwargs):
#   def weight_matrices(self):
# def calculate_linear_regressor_parameters_from_dataset(dataset, lamb=0.1,
#                                                        use_offset=True,
#                                                        use_ridge=True):
#   def __init__(self, input_dataset, num_hidden_list=None, **kwargs):
#   def call(self, input_dataset):
#   def compile(self, optimizer=tf.keras.optimizers.RMSprop,
#               loss='mse',
#               metrics=(pearson_correlation_first, 'mse'),
#               learning_rate=1e-3, **kwargs):
#   def fit(self, input_dataset, **kwargs):
#   def __init__(self, input_dataset, num_hidden_list=None, **kwargs):
#   def compile(self, optimizer=tf.keras.optimizers.Adam,
#               loss=tf.keras.losses.BinaryCrossentropy(),
#               metrics='accuracy', learning_rate=1e-3, **kwargs):
#   def call(self, input_dataset):
#   def fit(self, input_dataset, **kwargs):
# class PearsonCorrelationLoss(tf.keras.losses.Loss):
# class BrainModel(tf.keras.models.Model):
# class BrainModelLinearRegression(BrainModel):
# class BrainModelDNN(BrainModel):
# class BrainModelClassifier(BrainModel):


# Path: telluride_decoding/cca.py
from absl import logging
from telluride_decoding import brain_model
import numpy as np
import tensorflow.compat.v2 as tf

                                shape=(self.input1_dim, real_output_dims),
                                initializer=self._rot1_init,
                                trainable=False)
    self.rot2 = self.add_weight(name='rot2',
                                shape=(self.input2_dim, real_output_dims),
                                initializer=self._rot2_init,
                                trainable=False)
    logging.info('BrainCcaLayer.build mean1=%s, mean2=%s, rot1=%s, rot2=%s',
                 self.mean1.shape, self.mean2.shape,
                 self.rot1.shape, self.rot2.shape)

  def call(self, inputs):
    input1 = inputs[0]
    input2 = inputs[1]

    input1 = tf.reshape(input1, (-1, self.input1_dim))
    input2 = tf.reshape(input2, (-1, self.input2_dim))

    rotated_input1 = tf.keras.backend.dot(input1 - self.mean1, self.rot1)
    rotated_input2 = tf.keras.backend.dot(input2 - self.mean2, self.rot2)

    return tf.concat((rotated_input1, rotated_input2),
                     axis=1, name='ConcatenateCcaResults')

  def get_config(self):
    """Defines the parameters needed when re-creating this layer."""
    return {'requested_cca_dims': self.output_dims,
           }



class BrainModelCCA(brain_model.BrainModel):