import logging
import timeit as timeit
import os
import numpy as np

import pytest
from fontTools.mtiLib import parseSingleSubst

import psyneulink as pnl
from psyneulink import RANDOM_CONNECTIVITY_MATRIX

from psyneulink.core.components.functions.nonstateful.transferfunctions import Logistic
from psyneulink.core.components.functions.nonstateful.learningfunctions import BackPropagation
from psyneulink.core.compositions.composition import Composition
from psyneulink.core.globals import Context
from psyneulink.core.globals.keywords import Loss, DEFAULT_LEARNING_RATE, TRAINING_SET
from psyneulink.core.components.mechanisms.processing.transfermechanism import TransferMechanism
from psyneulink.core.components.projections.pathway.mappingprojection import MappingProjection
from psyneulink.library.compositions.autodiffcomposition import AutodiffComposition, AutodiffCompositionError
from psyneulink.core.compositions.report import ReportOutput


logger = logging.getLogger(__name__)


# All tests are set to run. If you need to skip certain tests,
# see http://doc.pytest.org/en/latest/skipping.html

# Unit tests for functions of AutodiffComposition class that are new (not in Composition)
# or override functions in Composition

def _single_learn_results(composition, *args, **kwargs):
    kwargs['synch_results_with_torch'] = 'run'
    composition.learn(*args, **kwargs)
    return composition.learning_results


TORCH_PARAM = 'torch_param'
PROJECTION = 'projection'
PROJ = 'proj'
NAME = 'proj_name'

@pytest.mark.pytorch
@pytest.mark.autodiff_constructor
class TestAutodiffConstructor:

    def test_no_args(self):
        comp = AutodiffComposition()
        assert isinstance(comp, AutodiffComposition)

    def test_two_calls_no_args(self):
        comp = AutodiffComposition()
        comp_2 = AutodiffComposition()
        assert isinstance(comp, AutodiffComposition)
        assert isinstance(comp_2, AutodiffComposition)

    # KAM removed this pytest 10/30 after removing target_CIM
    # def test_target_CIM(self):
    #     comp = AutodiffComposition()
    #     assert isinstance(comp.target_CIM, CompositionInterfaceMechanism)
    #     assert comp.target_CIM.composition == comp
    #     assert comp.target_CIM_ports == {}

    def test_no_initial_pytorch_representation(self):
        comp = AutodiffComposition()
        assert comp.pytorch_representation is None

    def test_duplicate_projections_to_nested_comp(self):
        input_node_autodiff = pnl.ProcessingMechanism(name='autodiff INPUT', input_shapes=2)
        hidden_node_autodiff_1 = pnl.ProcessingMechanism(name='autodiff HIDDEN 1', input_shapes=3)
        hidden_node_autodiff_2 = pnl.ProcessingMechanism(name='autodiff HIDDEN 2', input_shapes=4)
        hidden_node_autodiff_3 = pnl.ProcessingMechanism(name='autodiff HIDDEN 3', input_shapes=5)
        output_node_autodiff = pnl.ProcessingMechanism(name='autodiff OUTPUT', input_shapes=3)

        nested = pnl.AutodiffComposition(name='autodiff NESTED',
                                     nodes = [hidden_node_autodiff_1,
                                              hidden_node_autodiff_2])

        nested = pnl.AutodiffComposition(name='autodiff NESTED 2',
                                     nodes = [hidden_node_autodiff_1,
                                              hidden_node_autodiff_2])
        pathway_a = [input_node_autodiff, MappingProjection(input_node_autodiff, hidden_node_autodiff_1), nested]
        pathway_b = [input_node_autodiff, MappingProjection(input_node_autodiff, hidden_node_autodiff_2), nested,
                     MappingProjection(sender=hidden_node_autodiff_2, name="PROBLEM PROJ"), output_node_autodiff]
        autodiff_comp = pnl.AutodiffComposition(pathways=[pathway_a, pathway_b], name='autodiff COMP')
        with pytest.raises(AutodiffCompositionError) as error_text:
            autodiff_comp._build_pytorch_representation()
        assert error_text.value.error_value == ("First afferent Projection to 'autodiff HIDDEN 1' (which should be "
                                                "from 'autodiff NESTED Input_CIM') is not the same as its Projection "
                                                "from the input_CIM of 'autodiff NESTED 2'. One for this reason may "
                                                "be that these Components belong to different Compositions.")

    @pytest.fixture
    def copy_test_components(self):
        import torch
        def _get_copy_test_components(use_slice, proj_spec):
            if use_slice is True:
                # GRU uses sets of weights defined as slices within a single Parameter
                torch_module = torch.nn.GRU(input_size=3, hidden_size=5, bias=False)
                torch_parameter = list(torch_module.parameters())[0]
                torch_tensor = torch_module.state_dict()['weight_ih_l0'][slice(0,5)]
                autodiff = pnl.GRUComposition(input_size=3, hidden_size=5, bias=False)
                proj_name = 'INPUT TO NEW WEIGHTS'
                proj = autodiff.projections[proj_name]
                proj_spec = proj if proj_spec == PROJ else proj_name
                torch_param_specs = {#  name: proj_spec, torch_param, torch_module, slice, error_msg
                    'tensor_slice':(proj_spec, torch_module.state_dict()['weight_ih_l0'], None, slice(0,5), None),
                    'param_slice':(proj_spec, torch_parameter, None, slice(0,5), None),
                    'module_with_param_name_and_slice':(proj_spec, 'weight_ih_l0', torch_module, slice(0,5), None),
                    'shape_mismatch': (proj_spec, 'weight_hh_l0', torch_module,slice(0,5),
                                       f"Shape of torch parameter (5, 5) in copy_torch_param_to_projection_matrix() "
                                       f"does not match shape of matrix for 'INPUT TO NEW WEIGHTS' (3, 5)."),
                    'tensor_with_bad_slice':(proj_spec,torch_module.state_dict()['weight_ih_l0'],None, 'SLICE OF PI',
                                             "Specification of 'torch_slice' arg in "
                                             "copy_torch_param_to_projection_matrix() (SLICE OF PI) must be a slice."),
                    'param_with_bad_slice':(proj_spec, torch_parameter, None, 'SLICE OF PI',
                                            f"Specification of 'torch_slice' arg in "
                                            f"copy_torch_param_to_projection_matrix() (SLICE OF PI) must be a slice."),
                    'module_with_param_name_and_bad_slice':(proj_spec, 'weight_hh_l0', torch_module, 'SLICE OF PI',
                                                            f"Specification of 'torch_slice' arg in "
                                                            f"copy_torch_param_to_projection_matrix() ('SLICE OF PI') "
                                                            f"for Parameter 'weight_hh_l0' of GRU(3, 5, bias=False) "
                                                            f"must be a slice."),
                    'module_with_param_index_and_bad_slice':(proj_spec, 0, torch_module, 'SLICE OF PI',
                                                             f"Specification of 'torch_slice' arg in "
                                                             f"copy_torch_param_to_projection_matrix() "
                                                             f"('SLICE OF PI') for Parameter 0 of "
                                                             f"GRU(3, 5, bias=False) must be a slice.")
                }
            elif use_slice is False:
                # No slices
                torch_module = torch.nn.Linear(3, 5, bias=False)
                torch_parameter = list(torch_module.parameters())[0]
                torch_tensor = torch_module.state_dict()['weight']
                autodiff = pnl.AutodiffComposition([pnl.ProcessingMechanism(input_shapes=3),
                                                    pnl.MappingProjection(name='PROJECTION'),
                                                    pnl.ProcessingMechanism(input_shapes=5)])
                proj_name = 'PROJECTION'
                proj = autodiff.projections[proj_name]
                proj_spec = proj if proj_spec == PROJ else proj_name

                torch_param_specs = {#  name: proj_spec, torch_param, torch_module, slice, error_msg
                    'tensor': (proj_spec, torch_tensor, None, None, None),
                    'param': (proj_spec, torch_parameter, None, None, None),
                    'module_with_param_name': (proj_spec, 'weight', torch_module, None, None),
                    'module_with_param_index': (proj_spec, 0, torch_module, None, None),
                    'shape_mismatch': (proj_spec, torch.zeros(4), None, None,
                                       f"Shape of torch parameter (1, 4) in copy_torch_param_to_projection_matrix() "
                                       f"does not match shape of matrix for 'PROJECTION' (3, 5). [Note: torch biases, "
                                       f"usually 1d, have already been converted to 2d to match "
                                       f"PsyNeuLink BIAS Nodes Projections.]"),
                    'bad_param_spec': (proj_spec, set(), None, None,
                                       f"Specification of 'torch_param' arg in copy_torch_param_to_projection_matrix() "
                                       f"(set()) must be a torch.nn.Parameter, torch.Tensor, str or int."),
                    'module_in_param_spec': (proj_spec, torch_module, None, None,
                                             f"Specification of 'torch_param' arg in "
                                             f"copy_torch_param_to_projection_matrix() (Linear(in_features=3, "
                                             f"out_features=5, bias=False)) is a Module, but must be a "
                                             f"torch.nn.Parameter, torch.Tensor, str or int; if a Module is intended, "
                                             f"use the 'torch_module' arg, and specify the Parameter name or index in "
                                             f"the 'torch_param' arg."),
                    'param_str_without_module_spec': (proj_spec, "'STRING'", None, None,
                                                      f"Specifying of the 'torch_param' arg in "
                                                      f"copy_torch_param_to_projection_matrix() with a "
                                                      f"string or int ('STRING') requires the 'torch_module' arg "
                                                      f"to be specified as well."),
                    'param_int_without_module_spec': (proj_spec, 3, None, None,
                                                      f"Specifying of the 'torch_param' arg in "
                                                      f"copy_torch_param_to_projection_matrix() with a "
                                                      f"string or int (3) requires the 'torch_module' arg "
                                                      f"to be specified as well."),
                    'param_none_spec': (proj_spec, None, "'I DON'T MATTER'", None,
                                        f"The 'torch_param' arg in copy_torch_param_to_projection_matrix() (None) "
                                        f"must be specified, using either a torch.nn.Parameter or torch.Tensor, or a "
                                        f"str or int paired with specification of a torch.nn.Module in the "
                                        f"'torch_module' arg."),
                    'param_in_module_spec': (proj_spec, None, torch_parameter, None,
                                             f"Specification of 'torch_module' arg in "
                                             f"copy_torch_param_to_projection_matrix() is a torch Parameter or Tensor; "
                                             f"this should be specified using the 'torch_para' arg."),
                    'bad_module_spec': (proj_spec, 3, "'I BE BAD MODULE'", None,
                                        f"Specification of 'torch_module' arg in "
                                        f"copy_torch_param_to_projection_matrix() "
                                        f"('I BE BAD MODULE') must be a torch.nn.Module."),
                    'module_without_param_spec': (proj_spec, None, torch_module, None,
                                                  f"The 'torch_param' arg in copy_torch_param_to_projection_matrix() "
                                                  f"(None) must be specified, using either a torch.nn.Parameter or "
                                                  f"torch.Tensor, or a str or int paired with specification of a "
                                                  f"torch.nn.Module in the 'torch_module' arg."),
                    'param_not_in_state_dict': (proj_spec, "I'M IN A BAD STATE", torch_module, None,
                                                f"'I'M IN A BAD STATE' specified in 'torch_param' arg of "
                                                f"copy_torch_param_to_projection_matrix() is not the name of a "
                                                f"Parameter in the state_dict() for "
                                                f"'Linear(in_features=3, out_features=5, bias=False)'."),
                    'param_index_out_of_range': (proj_spec, 3, torch_module, None,
                                                 f"The value (3) specified in the 'torch_param' arg of "
                                                 f"copy_torch_param_to_projection_matrix() is not an index within "
                                                 f"the range of the ParameterList specified for the Module "
                                                 f"('Linear(in_features=3, out_features=5, bias=False)')."),
                    'bad_projection_name': ("BAD NAME", torch_tensor, None, None,
                                            f"'BAD NAME' in copy_torch_param_to_projection_matrix() "
                                            f"is not the name of a Projection in 'autodiff_composition'."),
                    'bad_projection': (MappingProjection(), torch_tensor, None, None,
                                       f"'Deferred Init MappingProjection' in copy_torch_param_to_projection_matrix() "
                                       f"is not a Projection in 'autodiff_composition'.")
                }
            else:
                assert False, f"Invalid use_slice value: {use_slice}"
            return torch_tensor, torch_param_specs, autodiff, proj
        return _get_copy_test_components

    # Test cases for copy_torch_param_to_projection_matrix()
    #                              (test, use_slice, proj_spec)
    copy_test_params = [
        # name, use_slice, validate, proj_spec
        # Valid specifications, with validation
        ('tensor_slice', True, True, NAME),
        ('param_slice', True, True, NAME),
        ('module_with_param_name', False, True, PROJ),
        ('module_with_param_index', False, True, NAME),
        ('module_with_param_name_and_slice', True, True, PROJ),
        # Valid specifications, without validation
        ('tensor', False, False, PROJ),
        ('param', False, False, PROJ),
        ('tensor_slice', True, False, PROJ),
        ('param_slice', True, False, PROJ),
        # Invalid specifications, without validation (should generate errors)
        ('shape_mismatch', False, True, NAME),
        ('shape_mismatch', True, True, PROJ),
        ('bad_param_spec', False, True, NAME),
        ('module_in_param_spec', False, True, PROJ),
        ('param_str_without_module_spec', False, True, NAME),
        ('param_int_without_module_spec', False, True, PROJ),
        ('param_none_spec', False, True, NAME),
        ('bad_module_spec', False, True, NAME),
        ('param_in_module_spec', False, True, PROJ),
        ('module_without_param_spec', False, True, NAME),
        ('param_not_in_state_dict', False, True, PROJ),
        ('param_index_out_of_range', False, True, NAME),
        ('bad_projection_name', False, True, PROJ),
        ('bad_projection', False, True, NAME),
        ('tensor_with_bad_slice', True, True, PROJ),
        ('param_with_bad_slice', True, True, NAME),
        ('module_with_param_name_and_bad_slice', True, True, PROJ),
        ('module_with_param_index_and_bad_slice', True, True, NAME),
    ]

    copy_test_method = ['torch_to_pnl', 'pnl_to_torch']

    @pytest.mark.parametrize('method', copy_test_method, ids=[x for x in copy_test_method])
    @pytest.mark.parametrize('test_condition, use_slice, validate, proj_spec', copy_test_params,
                             ids=[x[0] + ('_validate' if x[2] else '_no_validation')
                                  for x in copy_test_params])
    def test_torch_param_projection_matrix_exchange_methods(self,
                                                            test_condition,
                                                            use_slice,
                                                            validate,
                                                            proj_spec,
                                                            method,
                                                            copy_test_components):
        torch_tensor, torch_param_specs, autodiff, proj = copy_test_components(use_slice, proj_spec)
        proj_spec, torch_param, torch_module, torch_slice, error_msg = torch_param_specs[test_condition]

        if method == 'torch_to_pnl':
            copy_method = autodiff.copy_torch_param_to_projection_matrix
        else:
            copy_method = autodiff.copy_projection_matrix_to_torch_param

        if error_msg is None:
            copy_method(proj_spec, torch_param, torch_module, torch_slice, validate)
            torch_param_as_pnl_matrix = torch_tensor.detach().cpu().clone().numpy().T
            new_matrix = autodiff.projections[proj].parameters.matrix.get()
            np.testing.assert_allclose(new_matrix, torch_param_as_pnl_matrix)

        else:
            with pytest.raises(AutodiffCompositionError) as error_text:
                copy_method(proj_spec, torch_param, torch_module, torch_slice, validate)
            assert error_text.value.error_value == error_msg


    def test_report_prefs(self):
        comp = AutodiffComposition()
        assert comp.input_CIM.reportOutputPref == ReportOutput.OFF
        assert comp.output_CIM.reportOutputPref == ReportOutput.OFF
        # assert comp.target_CIM.reportOutputPref == False

    # FIXME: This test for patience doesn't actually test for correctness
    # def test_patience(self):
        # comp = AutodiffComposition()
        # assert comp.patience == 10


# Expected results for test_projection_specific_learning_rates()
# NOTE: these should be kept consistent with test_learning/test_projection_specific_learning_rates()
#       to additionally test for identicality of effects with Python learning
baseline = [[4.06551247, 4.06551247, 4.06551247, 4.06551247, 4.06551247]]
learn_method = [[0.03072, 0.03072, 0.03072, 0.03072, 0.03072]]
input_proj = [[1.0479138, 1.0479138, 1.0479138, 1.0479138, 1.0479138]]
hidden_proj = [[5.55952143, 5.55952143, 5.55952143, 5.55952143, 5.55952143]]
inpt_learn_ovrd = [[0.00768, 0.00768, 0.00768, 0.00768, 0.00768]]
hid_learn_ovrd = [[-0.49108492, -0.49108492, -0.49108492, -0.49108492, -0.49108492]]
input_dict_cnstr = [[3.2844555, 3.2844555, 3.2844555, 3.2844555, 3.2844555]]
hid_dict_constr = [[11.45824471, 11.45824471, 11.45824471, 11.45824471, 11.45824471]]
inp_cnstr_ovrd = [[0.12288, 0.12288, 0.12288, 0.12288, 0.12288]]
hid_cnstr_ovrd = [[-0.14778898, -0.14778898, -0.14778898, -0.14778898, -0.14778898]]
default_lr = .01
# [[1.51359819 1.51359819 1.51359819 1.51359819 1.51359819]]

class TestAutodiffLearningRateArgs:
    test_args = [
        # NOTES:
        #   Have to explicitly specify default_lr in constructor here (when it is expected to have an effect),
        #      since default learning_rates are different for Composition (.05) and  AutodiffComposition (.001)
        #   Tests both Projection *to* nested Composition (input) and Projection *in* nested Comopsition (hidden)
        #   condition          constructor_lr   lrn()_lr  inp_lr  hid_lr  cstr_dict learn_dict post_cstr expected
        ("baseline",               default_lr,    None,    None,    None,    None,     None,    False,   baseline),
        # learn() lr
        ("lrn_method",              None,          .1,     None,    None,    None,     None,    False,   learn_method),
        # learn() lr overrides constructor lr
        ("lrn_override",           default_lr,     .1,     None,    None,    None,     None,    False,   learn_method),
        # Projection lr (in constructor)
        ("input_proj",             default_lr,    None,    .2,      None,    None,     None,    False,   input_proj),
        ("hidden_proj",            default_lr,    None,    None,     .2,     None,     None,    False,   hidden_proj),
        # Projection lr overrides specificdation of default lr in learning()
        #   note: learn_dict = None -> don't include either Projection specification in learn();
        #         assignments of None to entries in the dict itself are tested in
        #         test_projection_specific_learning_rates() and test_learning/test_default_and_False_learning_rates()
        ("inpt_override_lrn",      default_lr,     .1,     .2,      None,    None,     None,    False,   inpt_learn_ovrd),
        ("hidn_override_lrn",      default_lr,     .1,     None,     .2,     None,     None,    False,   hid_learn_ovrd),
        # # Specifications of learning_rates in dict in Autodiff constructor overrides Projection lr
        ("input_dict_constructor", default_lr,    None,     .2,     None,  'input',    None,    False,   input_dict_cnstr),
        ("hidden_dict_constructor",default_lr,    None,    None,     .2,   'hidden',   None,    False,   hid_dict_constr),
        # learning_rate dict in learn() overrides Projection lr
        ("input_dict_learn",       default_lr,    None,     .3,     None,   None,    'input',   False,   input_proj),
        ("hidden_dict_learn",      default_lr,    None,    None,     .3,    None,    'hidden',  False,   hidden_proj),
        # Projection spec in learning_rate dict in Autodiff constructor overrides default lr and Projection lr
        ("inpt_ovrd_constructor",  default_lr,     .1,      .2,     None,  'input',    None,    False,  inp_cnstr_ovrd),
        ("hid_ovrd_constructor",   default_lr,     .1,     None,     .2,   'hidden',   None,    False,  hid_cnstr_ovrd),
        # Projection spec in learning_rate dict in learn() overrides default lr and Projection lr
        #                                (use two different Projection lr's to be sure they don't matter)
        ("in_ovrd_learn_method3",  default_lr,     .1,      .3,     None,   None,    'input',   False, inpt_learn_ovrd),
        ("hid_ovrd_learn_method3", default_lr,     .1,     None,     .3,    None,    'hidden',  False,  hid_learn_ovrd),
        ("in_ovrd_learn_method5",  default_lr,     .1,      .5,     None,   None,    'input',   False, inpt_learn_ovrd),
        ("hid_ovrd_learn_method5", default_lr,     .1,     None,     .5,    None,    'hidden',  False,  hid_learn_ovrd),
        # Projection spec in learning_rate dict in learn() overrides everything else, including constructor dict
        ("inpt_learn_constructor", default_lr,     .1,      .4,     None,  'input',  'input',   False,   inpt_learn_ovrd),
        ("hid_learn_constructor",  default_lr,     .1,     None,     .4,   'hidden', 'hidden',  False,   hid_learn_ovrd),
        # Projection spec made after AutodiffComposition contruction (should have no effect, since post-construction)
        ("inpt_learn_constructor", default_lr,     .1,      .4,     None,  'input',  'input',   True,    inpt_learn_ovrd),
        ("hid_learn_constructor",  default_lr,     .1,     None,     .4,   'hidden', 'hidden',  True,    hid_learn_ovrd),
    ]
    # NOTE: this should be kept consistent with test_learning/test_projection_specific_learning_rates()
    #       to additionally test for identicality of effects with Python learning
    @pytest.mark.pytorch
    @pytest.mark.composition
    @pytest.mark.parametrize("condition, constructor_lr, learn_method_lr, "
                             "input_lr, hidden_lr, constructor_dict_param, learn_dict_param, post_constr, expected",
                             test_args, ids=[f"{x[0]}" for x in test_args])
    def test_projection_specific_learning_rates(self, condition, constructor_lr, learn_method_lr, input_lr, hidden_lr,
                                                constructor_dict_param, learn_dict_param, post_constr, expected):
        in_shape = 4
        nested_1_shape = 3
        nested_2_shape = 2
        out_shape = 5
        input_stims = [[.1,.2,.3,.4]]
        target_vals = [[1,1,1,1,1]]
        num_trials = 3

        # Construct nested_comp
        nested_mech_1 = pnl.ProcessingMechanism(name='Nested Mech 1', input_shapes=nested_1_shape)
        nested_mech_2 = pnl.ProcessingMechanism(name='Nested Mech 2', input_shapes=nested_2_shape)
        nested_proj = pnl.MappingProjection(nested_mech_1, nested_mech_2, name="INNER PROJECTION",
                                            learning_rate=hidden_lr)
        nested_comp = pnl.AutodiffComposition(name='Inner Comp',
                                              pathways=[nested_mech_1, nested_proj, nested_mech_2],
                                              # Need to explicitly set this to maintain compatibility with PNL
                                              learning_rate = default_lr)

        outer_mech_in = pnl.ProcessingMechanism(name='Outer Mech IN', input_shapes=in_shape)
        outer_mech_out = pnl.ProcessingMechanism(name='Outer Mech OUT', input_shapes=out_shape)
        input_proj = pnl.MappingProjection(outer_mech_in, nested_mech_1,
                                           name="OUTER INPUT PROJECTION",
                                           learning_rate=input_lr)
        pathway = [outer_mech_in, input_proj, nested_comp, outer_mech_out]
        constructor_learning_rate_dict = {input_proj: .3 if constructor_dict_param == 'input' else None,
                                          nested_proj: .3 if constructor_dict_param == 'hidden' else None,
                                          pnl.DEFAULT_LEARNING_RATE: constructor_lr}

        learn_method_learning_rate_dict = {pnl.DEFAULT_LEARNING_RATE: learn_method_lr}
        if learn_dict_param == 'input':
            learn_method_learning_rate_dict.update({input_proj: .2})
        if learn_dict_param == 'hidden':
            learn_method_learning_rate_dict.update({nested_proj: .2})

        # Construct outer_comp
        outer_comp = pnl.AutodiffComposition(name='Outer Comp',
                                             pathways=pathway,
                                             learning_rate = (constructor_learning_rate_dict
                                                              if "constructor" in condition
                                                              else {pnl.DEFAULT_LEARNING_RATE: constructor_lr}))
        # Test for specified learning rates in constructor
        input_proj_lr = \
            .3 if constructor_dict_param == 'input' \
                else input_lr if input_lr \
                else constructor_lr if constructor_lr is not None \
                else .001
        hidden_proj_lr = \
            .3 if constructor_dict_param == 'hidden' \
                else hidden_lr if hidden_lr \
                else constructor_lr if constructor_lr is not None \
                else default_lr
        outer_comp_proj_lr = constructor_lr if constructor_lr is not None else .001
        pytorch_rep = outer_comp._build_pytorch_representation()
        assert pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == input_proj_lr # (vs. .2 in "...lr_2_MOD"
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_proj) == hidden_proj_lr
        assert pytorch_rep.get_torch_learning_rate_for_projection(outer_comp.projections[1]) == outer_comp_proj_lr

        # This should have no effet since done post construction
        if post_constr:
            input_proj.learning_rate = .9
            nested_proj.learning_rate = .7

        # Test learning
        targets = outer_comp.infer_backpropagation_learning_pathways(pnl.ExecutionMode.PyTorch)
        pytorch_result = outer_comp.learn(
            inputs={outer_mech_in:input_stims, targets[0]: target_vals},
            num_trials=num_trials,
            execution_mode=pnl.ExecutionMode.PyTorch,
            # learning_rate=learn_method_lr
            learning_rate=(learn_method_learning_rate_dict if "learn" in condition
                           else {pnl.DEFAULT_LEARNING_RATE: learn_method_lr}))

        pytorch_rep = outer_comp.parameters.pytorch_representation.get('Outer Comp')
        # assert pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == 0.3 # (vs. .2 in "...lr_2_MOD"
        assert (pytorch_rep.get_torch_learning_rate_for_projection(input_proj) ==
                (learn_method_learning_rate_dict[input_proj] if learn_dict_param == 'input'
                 else constructor_learning_rate_dict[input_proj] if constructor_dict_param == 'input'
                else input_lr if input_lr not in {None, True}
                # else learn_method_lr if (pnl.DEFAULT_LEARNING_RATE in learn_method_learning_rate_dict and
                #                          learn_method_learning_rate_dict[pnl.DEFAULT_LEARNING_RATE] not in {None, True})
                else learn_method_lr if (learn_method_learning_rate_dict and learn_method_lr not in {None, True})
                else .9 if post_constr
                else default_lr))
        assert (pytorch_rep.get_torch_learning_rate_for_projection(nested_proj) ==
                (learn_method_learning_rate_dict[nested_proj] if learn_dict_param == 'hidden'
                 else constructor_learning_rate_dict[nested_proj] if constructor_dict_param == 'hidden'
                else hidden_lr if hidden_lr not in {None, True}
                else learn_method_lr if (learn_method_learning_rate_dict and learn_method_lr not in {None, True})
                else .7 if post_constr
                else default_lr))
        assert (pytorch_rep.get_torch_learning_rate_for_projection(outer_comp.projections[1]) ==
                (learn_method_lr if learn_method_lr is not None
                 else constructor_lr if constructor_lr is not None
                else .1))

        np.testing.assert_allclose(pytorch_result, expected)

    @pytest.mark.pytorch
    @pytest.mark.parametrize("build_pytorch_rep_spec", [None, 10, 'dict'])
    def test_learning_rate_with_3_levels_of_nesting_and_build_pytorch_rep_spec(self, build_pytorch_rep_spec):
        # Test construction with enable_learning=False, running, and then enabling learning
        if build_pytorch_rep_spec == 'dict':
            build_pytorch_rep_spec = {DEFAULT_LEARNING_RATE: 10,
                                      "NESTED 2 PROJ CD": .9 }
        nested_2_proj_AB_lr = .1
        nested_2_proj_BC_lr = None
        nested_2_proj_CD_lr = True
        nested_2_proj_DE_lr = False
        nested_2_comp_lr = None
        nested_1_proj_in_lr = .2
        nested_1_proj_out_lr = True
        nested_1_comp_lr = None if build_pytorch_rep_spec else .3
        input_proj_lr = True
        output_proj_lr = .4
        outer_comp_lr = None if build_pytorch_rep_spec else .5

        nested_2_mech_A = pnl.ProcessingMechanism(name='NESTED 2 NODE A', input_shapes=2)
        nested_2_mech_B = pnl.ProcessingMechanism(name='NESTED 2 NODE B', input_shapes=5)
        nested_2_mech_C = pnl.ProcessingMechanism(name='NESTED 2 NODE C', input_shapes=4)
        nested_2_mech_D = pnl.ProcessingMechanism(name='NESTED 2 NODE D', input_shapes=3)
        nested_2_mech_E = pnl.ProcessingMechanism(name='NESTED 2 NODE E', input_shapes=2)
        nested_2_proj_AB = pnl.MappingProjection(nested_2_mech_A, nested_2_mech_B,
                                                 name="NESTED 2 PROJ AB",
                                                 matrix=RANDOM_CONNECTIVITY_MATRIX,
                                                 learning_rate=nested_2_proj_AB_lr)
        nested_2_proj_BC = pnl.MappingProjection(nested_2_mech_B, nested_2_mech_C,
                                                 name="NESTED 2 PROJ BC",
                                                 learning_rate=nested_2_proj_BC_lr)
        nested_2_proj_CD = pnl.MappingProjection(nested_2_mech_C, nested_2_mech_D,
                                                 name="NESTED 2 PROJ CD",
                                                 learning_rate=nested_2_proj_CD_lr)
        nested_2_proj_DE = pnl.MappingProjection(nested_2_mech_D, nested_2_mech_E,
                                                 name="NESTED 2 PROJ DE",
                                                 learning_rate=nested_2_proj_DE_lr)
        nested_2_comp = pnl.AutodiffComposition(name='Nested 2 Comp',
                                                pathways=[nested_2_mech_A, nested_2_proj_AB,
                                                          nested_2_mech_B, nested_2_proj_BC,
                                                          nested_2_mech_C, nested_2_proj_CD,
                                                          nested_2_mech_D, nested_2_proj_DE,
                                                          nested_2_mech_E],
                                                learning_rate=nested_2_comp_lr)

        nested_1_mech_IN = pnl.ProcessingMechanism(name='NESTED 1 NODE IN', input_shapes=4)
        nested_1_mech_OUT = pnl.ProcessingMechanism(name='NESTED 1 NODE OUT', input_shapes=3)
        nested_1_proj_in = pnl.MappingProjection(nested_1_mech_IN, nested_2_mech_A,
                                                 name="NESTED 1 PROJ IN",
                                                 learning_rate=nested_1_proj_in_lr)
        nested_1_proj_out = pnl.MappingProjection(nested_2_mech_E, nested_1_mech_OUT,
                                                  learning_rate=nested_1_proj_out_lr,
                                                  name="NESTED 1 PROJ OUT")
        nested_1_comp = pnl.AutodiffComposition(name='Nested 1 Comp',
                                                pathways=[nested_1_mech_IN,
                                                          nested_1_proj_in,
                                                          nested_2_comp,
                                                          nested_1_proj_out,
                                                          nested_1_mech_OUT],
                                                learning_rate = nested_1_comp_lr)

        outer_mech_in = pnl.ProcessingMechanism(name='OUTER NODE IN')
        outer_mech_out = pnl.ProcessingMechanism(name='OUTER NODE OUT')
        input_proj = pnl.MappingProjection(outer_mech_in, nested_1_mech_IN, learning_rate=input_proj_lr,
                                           name="OUTER PROJ TO NESTED 1")
        output_proj = pnl.MappingProjection(nested_1_mech_OUT, outer_mech_out, learning_rate=output_proj_lr,
                                            name="OUTER PROJ FROM NESTED 1")
        outer_comp = pnl.AutodiffComposition([outer_mech_in, input_proj, nested_1_comp, output_proj, outer_mech_out],
                                             name='Outer Comp',
                                             learning_rate = outer_comp_lr)

        outer_comp._build_pytorch_representation(build_pytorch_rep_spec)

        # learning_rates should reflect values specified in _build_pytorch_repreentation() and learn()
        outer_comp.learn(inputs={outer_mech_in: [[1]],
                                 outer_comp.get_target_nodes()[0]: [[1]]},
                         num_trials=2, execution_mode=pnl.ExecutionMode.PyTorch,
                         learning_rate={input_proj:12,
                                        "NESTED 2 PROJ BC": 13,
                                        nested_2_proj_DE: 15})
        pytorch_rep = outer_comp.parameters.pytorch_representation.get('Outer Comp')
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_AB) == nested_2_proj_AB_lr
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_BC) == 13
        if build_pytorch_rep_spec:
            assert (pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_CD) ==
                    build_pytorch_rep_spec["NESTED 2 PROJ CD"] if isinstance(build_pytorch_rep_spec, dict) else 10)
        else:
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_CD) == nested_1_comp_lr
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_DE) == 15
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_1_proj_in) == nested_1_proj_in_lr
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_1_proj_out) == (nested_1_comp_lr or 10)
        assert pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == 12
        assert pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == output_proj_lr

        outer_comp.learn(inputs={outer_mech_in: [[1]], outer_comp.get_target_nodes()[0]: [[1]]},
                         num_trials=2, execution_mode=pnl.ExecutionMode.PyTorch)

        # learning_rates should revert to constructor-speficied values (e.g., input_proj) or composition defaults
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_AB) == nested_2_proj_AB_lr
        if build_pytorch_rep_spec:
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_BC) == 10
            assert (pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_CD) ==
                    build_pytorch_rep_spec["NESTED 2 PROJ CD"] if isinstance(build_pytorch_rep_spec, dict) else 10)
        else:
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_BC) == nested_1_comp_lr
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_CD) == nested_1_comp_lr
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_DE) is False
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_1_proj_in) == nested_1_proj_in_lr
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_1_proj_out) == (nested_1_comp_lr or 10)
        assert pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == (outer_comp_lr or 10)
        assert pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == output_proj_lr

        # Test additional calls to _build_pytorch_representation()
        if build_pytorch_rep_spec is None:  # just need to do these tests once:

            # check that nothing happens without specifying new in the call
            with pytest.warns(UserWarning) as warning:
                outer_comp._build_pytorch_representation(build_pytorch_rep_spec)
                assert ("The '_build_pytorch_representation() method for 'Outer Comp' has already been called "
                        "directly from the command line; this and any additional calls will be ignored. Make any "
                        "desired modifications to parameters (e.g., learning_rates) either in the constructor for "
                        "the AutodiffComposition, or its learn() method. in warning[0].message.args[0]")

            # change a projection learning_rate for composition using another call to _build_pytorch_representation()
            pytorch_rep = outer_comp._build_pytorch_representation(learning_rate={"NESTED 2 PROJ CD": 14})
            # check that it has taken effect:
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_BC) == .3
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_CD) == 14

            # check that it has been assigned as default (in _optimizer_constructor_params);
            #     i.e., that it persists after a call to learn():
            outer_comp.learn(inputs={outer_mech_in: [[1]], outer_comp.get_target_nodes()[0]: [[1]]},
                             learning_rate={"NESTED 2 PROJ BC": 99},)
            pytorch_rep_outer_comp = outer_comp.parameters.pytorch_representation.get('Outer Comp')
            assert pytorch_rep_outer_comp.get_torch_learning_rate_for_projection(nested_2_proj_BC) == 99
            assert pytorch_rep_outer_comp.get_torch_learning_rate_for_projection(nested_2_proj_CD) == 14
            pytorch_rep_constructor = outer_comp.parameters.pytorch_representation.get(None)
            assert pytorch_rep_constructor.get_torch_learning_rate_for_projection(nested_2_proj_CD) == 14

            # check that it persists after another call to learn() (vs. learning_rate specifed in the call to learn()
            outer_comp.learn(inputs={outer_mech_in: [[1]], outer_comp.get_target_nodes()[0]: [[1]]})
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_BC) == .3
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_2_proj_CD) == 14


    error_test_args = [
        ("comp_lr_spec_str",
         "A value ('hello') specified in the 'learning_rate' arg of the learn() method for 'Outer Comp' "
         "is not valid; it must be an int, float, bool or None."),
        ("comp_lr_spec_proj",
         "A value ('(MappingProjection INPUT PROJECTION)') specified in the 'learning_rate' arg of the learn() method "
         "for 'Outer Comp' is not valid; it must be an int, float, bool or None."),
        ("dict_lr_val_str",
         "A value ('goodbye') specified in the 'learning_rate' arg of the learn() method for 'Outer Comp' "
         "is not valid; it must be an int, float, bool or None."),
        ("dict_lr_val_proj",
         "A value ('(MappingProjection INPUT PROJECTION)') specified in the 'learning_rate' arg of the learn() method "
         "for 'Outer Comp' is not valid; it must be an int, float, bool or None."),
        ("dict_illegal_key_str",
         "The following Projection specified in the 'learning_rate' arg of the learn() method for 'Outer Comp' "
         "is not in that Composition or any nested within it: 'woa a woa'."),
        ("dict_illegal_key_int",
         "The following Projection specified in the 'learning_rate' arg of the learn() method for 'Outer Comp' "
         "is not in that Composition or any nested within it: '23'."),
        ("dict_key_bad_proj",
         "The following Projection specified in the 'learning_rate' arg of the learn() method for 'Outer Comp' "
         "is not in that Composition or any nested within it: 'BAD PROJECTION'."),
        ("dict_proj_not_learnable",
         "Projection ('INPUT PROJECTION') specified in the dict for the 'learning_rate' arg of the learn() method for "
         "'Outer Comp' is not learnable; check that its 'learnable' attribute is set to 'True' and its learning_rate "
         "is not 'False', or remove it from the dict.")
         ]
    @pytest.mark.pytorch
    @pytest.mark.parametrize("condition, error_msg", error_test_args,
                             ids=[f"{x[0]}" for x in error_test_args])
    def test_learning_rate_specification_errors(self, condition, error_msg):
        # Test for errors with learning_rates specified in Composition constructor
        nested_mech_1 = pnl.ProcessingMechanism(name='NESTED NODE 1')
        nested_mech_2 = pnl.ProcessingMechanism(name='NESTED NODE 2')
        nested_proj = pnl.MappingProjection(nested_mech_1, nested_mech_2, learning_rate=.3, name="NESTED PROJECTION")
        nested_comp = pnl.AutodiffComposition(name='Nested Comp', pathways=[nested_mech_1, nested_proj, nested_mech_2],)

        outer_mech_in = pnl.ProcessingMechanism(name='INPUT NODE')
        outer_mech_out = pnl.ProcessingMechanism(name='OUTPUT NODE')
        input_proj = pnl.MappingProjection(outer_mech_in, nested_mech_1,
                                           learning_rate=False if condition == 'dict_proj_not_learnable' else .2,
                                           name="INPUT PROJECTION")
        pathway = [outer_mech_in, input_proj, nested_comp, outer_mech_out]

        comp_lr = None
        default_lr = .1
        key_spec = input_proj
        val_spec = .2

        error_type = AutodiffCompositionError
        if condition == 'comp_lr_spec_str':
            comp_lr = 'hello'
            error_type = pnl.CompositionError
        elif condition == 'comp_lr_spec_proj':
            comp_lr = input_proj
            error_type = pnl.CompositionError
        elif condition == "dict_lr_val_str":
            val_spec = "goodbye"
        elif condition == "dict_lr_val_proj":
            val_spec = input_proj
        elif condition == "dict_illegal_key_str":
            key_spec = "woa a woa"
        elif condition == "dict_illegal_key_int":
            key_spec = 23
        elif condition == "dict_key_bad_proj":
            key_spec = pnl.MappingProjection(nested_mech_2, outer_mech_out, learning_rate=.4, name="BAD PROJECTION")
        elif condition == "dict_proj_not_learnable":
            error_type = AutodiffCompositionError
            comp_lr = None
            input_proj.learnable = False

        comp_lr = comp_lr or {DEFAULT_LEARNING_RATE: default_lr, key_spec: val_spec}

        with pytest.raises(error_type) as error_text:
            outer_comp = pnl.AutodiffComposition(pathway, name='Outer Comp')
            outer_comp.learn(inputs={outer_mech_in: [[1.0]]}, learning_rate=comp_lr)
        assert error_msg in str(error_text.value)

    @pytest.mark.pytorch
    def test_learning_rate_utility_functions(self):
        import torch
        nested_mech_1 = pnl.ProcessingMechanism(name='NESTED NODE 1')
        nested_mech_2 = pnl.ProcessingMechanism(name='NESTED NODE 2')
        nested_proj = pnl.MappingProjection(nested_mech_1, nested_mech_2, learning_rate=.3, name="NESTED PROJECTION")
        nested_comp = pnl.AutodiffComposition(name='Nested Comp', pathways=[nested_mech_1, nested_proj, nested_mech_2],)

        outer_mech_in = pnl.ProcessingMechanism(name='INPUT NODE')
        outer_mech_out = pnl.ProcessingMechanism(name='OUTPUT NODE')
        input_proj = pnl.MappingProjection(outer_mech_in, nested_mech_1,
                                           learning_rate=2.5,
                                           name="INPUT PROJECTION")
        output_proj = pnl.MappingProjection(nested_mech_2, outer_mech_out,
                                            learning_rate=.13,
                                            name="OUTPUT PROJECTION")
        outer_comp = pnl.AutodiffComposition([outer_mech_in, input_proj, nested_comp, output_proj, outer_mech_out],
                                             name='Outer Comp')

        pytorch_rep = outer_comp._build_pytorch_representation()
        assert pytorch_rep.get_torch_learning_rate_for_projection("INPUT PROJECTION") == 2.5
        assert pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == 2.5
        assert pytorch_rep.get_torch_learning_rate_for_projection('NESTED PROJECTION') == .3
        assert pytorch_rep.get_torch_learning_rate_for_projection(nested_proj) == .3
        assert pytorch_rep.get_torch_learning_rate_for_projection("OUTPUT PROJECTION") == .13
        assert pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == .13

        param_to_proj_dict = pytorch_rep.torch_params_to_projections()
        assert len(param_to_proj_dict) == 3
        assert all(isinstance(param, torch.nn.Parameter) for param in param_to_proj_dict.keys())
        assert all(isinstance(proj, MappingProjection) for proj in param_to_proj_dict.values())

        proj_to_param_dict = pytorch_rep.projections_to_torch_params()
        assert len(proj_to_param_dict) == 3
        assert all(isinstance(proj, MappingProjection) for proj in proj_to_param_dict.keys())
        assert all(isinstance(param, torch.nn.Parameter) for param in proj_to_param_dict.values())

        assert set(param_to_proj_dict.values()) == set(proj_to_param_dict.keys())
        assert set(proj_to_param_dict.values()) == set(param_to_proj_dict.keys())

    @pytest.mark.pytorch
    @pytest.mark.parametrize("enable_learning", ['before_run', 'after_run'])
    def test_enable_learning(self, enable_learning):
        # Test construction with enable_learning=False, running, and then enabling learning

        enable_learning = enable_learning == 'before_run'
        nested_proj_lr = .25
        input_proj_lr = True
        output_proj_lr = 1.9
        outer_comp_lr = .3

        nested_mech_1 = pnl.ProcessingMechanism(name='NESTED NODE 1', input_shapes=2)
        nested_mech_2 = pnl.ProcessingMechanism(name='NESTED NODE 2', input_shapes=3)
        nested_proj = pnl.MappingProjection(nested_mech_1, nested_mech_2,
                                            matrix=RANDOM_CONNECTIVITY_MATRIX,
                                            learning_rate=nested_proj_lr, name="NESTED PROJECTION")
        nested_comp = pnl.AutodiffComposition(name='Nested Comp', pathways=[nested_mech_1, nested_proj, nested_mech_2],
                                              enable_learning = enable_learning)

        outer_mech_in = pnl.ProcessingMechanism(name='INPUT NODE')
        outer_mech_out = pnl.ProcessingMechanism(name='OUTPUT NODE')
        input_proj = pnl.MappingProjection(outer_mech_in, nested_mech_1, learning_rate=input_proj_lr,
                                           name="INPUT PROJECTION")
        output_proj = pnl.MappingProjection(nested_mech_2, outer_mech_out, learning_rate=output_proj_lr,
                                            name="OUTPUT PROJECTION")
        outer_comp = pnl.AutodiffComposition([outer_mech_in, input_proj, nested_comp, output_proj, outer_mech_out],
                                             name='Outer Comp',
                                             enable_learning = enable_learning,
                                             learning_rate = outer_comp_lr)

        if enable_learning:
            # Execute learning without run
            learning_result = outer_comp.learn(inputs={outer_mech_in: [[1]], outer_comp.get_target_nodes()[0]: [[1]]},
                                               num_trials=2, execution_mode=pnl.ExecutionMode.PyTorch,
                                               learning_rate={input_proj:input_proj_lr})

        else:
            # Execute run, then enable learning and execute learning()
            run_result = outer_comp.run(inputs={outer_mech_in: [[1]], outer_comp.get_target_nodes()[0]: [[1]]},
                                        num_trials=2, execution_mode=pnl.ExecutionMode.PyTorch)

            # Check that learning rates are ones from construction
            # pytorch_rep = outer_comp._build_pytorch_representation(learning_rate=10)
            nested_comp.enable_learning = True
            outer_comp.enable_learning = True
            pytorch_rep = outer_comp._build_pytorch_representation()
            assert pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == outer_comp_lr
            assert pytorch_rep.get_torch_learning_rate_for_projection(nested_proj) == nested_proj_lr
            assert pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == output_proj_lr
            learning_result = outer_comp.learn(inputs={outer_mech_in: [[1]], outer_comp.get_target_nodes()[0]: [[1]]},
                                               num_trials=2, execution_mode=pnl.ExecutionMode.PyTorch,
                                               learning_rate={input_proj:input_proj_lr})
            # Ensure learning occurred
            assert learning_result != run_result

        # Should get same result after learning in both cases
        np.testing.assert_allclose(learning_result,[[-0.35035038561297505]])

# Note: don't use @pytest.mark.composition here to force test of Autodiff without torch installed
@pytest.mark.composition
def test_autodiff_without_torch():
    """Test Autodiff without torch installed"""

    # Should run with ExecutionMode.Python no matter what
    mech_A = TransferMechanism()
    mech_B = TransferMechanism()
    comp = AutodiffComposition([mech_A, mech_B])
    comp.run()
    result = comp.learn(inputs={mech_A:[[1], [2]]}, epochs=3, execution_mode=pnl.ExecutionMode.Python)

    assert comp.pytorch_representation is None
    np.testing.assert_allclose(result,[[1.95634283]], atol=1e-08, rtol=1e-08)

    # If torch is not installed, should raise exception for learn() with ExecutionMode.PyTorch
    from psyneulink.library.compositions.autodiffcomposition import torch_available
    if not torch_available:
        mech_A = TransferMechanism()
        mech_B = TransferMechanism()
        comp = AutodiffComposition([mech_A, mech_B])
        comp.run()
        with pytest.raises(AutodiffCompositionError) as error:
            result = comp.learn(inputs={mech_A:[[1], [2]]},
                                epochs=3,
                                execution_mode=pnl.ExecutionMode.PyTorch)
            np.testing.assert_allclose(result,[[1.95634283]], atol=1e-08, rtol=1e-08)
        assert (f"'autodiff_composition-1.learn()' has been called with ExecutionMode.Pytorch, "
                f"but Pytorch module ('torch') is not installed. "
                f"Please install it with `pip install torch` or `pip3 install torch`") in str(error.value)

@pytest.mark.pytorch
@pytest.mark.composition
def test_autodiff_forward(autodiff_mode):
    # create xor model mechanisms and projections
    xor_in = TransferMechanism(name='xor_in',
                               default_variable=np.zeros(2))

    xor_hid = TransferMechanism(name='xor_hid',
                                default_variable=np.zeros(10),
                                function=Logistic())

    xor_out = TransferMechanism(name='xor_out',
                                default_variable=np.zeros(1),
                                function=Logistic())

    hid_map = MappingProjection(matrix=np.random.rand(2,10))
    out_map = MappingProjection(matrix=np.random.rand(10,1))

    # put the mechanisms and projections together in an autodiff composition (AC)
    xor = AutodiffComposition()

    xor.add_node(xor_in)
    xor.add_node(xor_hid)
    xor.add_node(xor_out)

    xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
    xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

    outputs = xor.run(inputs=[0,0], execution_mode=autodiff_mode)
    np.testing.assert_allclose(outputs, [[0.9479085241082691]])

@pytest.mark.pytorch
@pytest.mark.composition
def test_retain_results():
    """Test that results from calls to learning() are added to results list (from retained_results)"""
    inputs = [[-0.8104468, -0.40517032, 0.75040168]]
    input_mech = pnl.ProcessingMechanism(input_shapes=3)
    output_mech = pnl.ProcessingMechanism(input_shapes=3)
    input_node = input_mech
    comp = AutodiffComposition([input_mech, output_mech])
    comp.run(inputs={input_node:inputs}, num_trials=1)
    comp.learn(inputs={input_node:inputs},num_trials=2)
    comp.run(inputs={input_node:inputs}, num_trials=3)
    comp.learn(inputs={input_node:inputs},num_trials=4)
    assert len(comp.results) == 10


@pytest.mark.pytorch
@pytest.mark.accorrectness
@pytest.mark.composition
class TestTrainingCorrectness:

    # test whether xor model created as autodiff composition learns properly
    @pytest.mark.benchmark(group="XOR")
    @pytest.mark.parametrize(
        'eps, calls, opt, expected', [
            (100, 'single', 'adam', [[[0.09823965]], [[0.81092879]], [[0.78179557]], [[0.25593583]]]),
            (50, 'multiple', 'adam', [[[0.31200036]], [[0.59406178]], [[0.60417587]], [[0.52347365]]]),
        ]
    )
    def test_xor_training_correctness(self, eps, calls, opt, autodiff_mode, benchmark, expected):
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection(matrix=np.random.rand(2, 10))
        out_map = MappingProjection(matrix=np.random.rand(10, 1))

        xor = AutodiffComposition(optimizer_type=opt, learning_rate=0.1)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        input_dict = {"inputs": {xor_in: np.array([[0, 0], [0, 1], [1, 0], [1, 1]])},
                      "targets": {xor_out: np.array([[0], [1], [1], [0]])}}

        if calls == 'single':
            input_dict["epochs"] = eps
            results = benchmark(_single_learn_results, xor, inputs=input_dict, execution_mode=autodiff_mode)
        else:
            input_dict["epochs"] = 1
            for i in range(eps - 1):
                xor.learn(inputs=input_dict, execution_mode=autodiff_mode)
            results = benchmark(_single_learn_results, xor, inputs=input_dict, execution_mode=autodiff_mode)

        if pytest.helpers.llvm_current_fp_precision() == 'fp32' and autodiff_mode != pnl.ExecutionMode.PyTorch:
            accuracy_args = {"atol": 1e-8, "rtol": 1e-6}
        else:
            accuracy_args = {}

        np.testing.assert_allclose(results, expected, **accuracy_args)


    # tests whether semantic network created as autodiff composition learns properly
    @pytest.mark.benchmark(group="Semantic net")
    @pytest.mark.parametrize(
        'eps, opt', [
            (50, 'adam'),
        ]
    )
    def test_semantic_net_training_correctness(self, eps, opt, autodiff_mode, benchmark):

        # MECHANISMS FOR SEMANTIC NET:

        nouns_in = TransferMechanism(name="nouns_input",
                                     default_variable=np.zeros(8))

        rels_in = TransferMechanism(name="rels_input",
                                    default_variable=np.zeros(3))

        h1 = TransferMechanism(name="hidden_nouns",
                               default_variable=np.zeros(8),
                               function=Logistic())

        h2 = TransferMechanism(name="hidden_mixed",
                               default_variable=np.zeros(15),
                               function=Logistic())

        out_sig_I = TransferMechanism(name="sig_outs_I",
                                      default_variable=np.zeros(8),
                                      function=Logistic())

        out_sig_is = TransferMechanism(name="sig_outs_is",
                                       default_variable=np.zeros(12),
                                       function=Logistic())

        out_sig_has = TransferMechanism(name="sig_outs_has",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        out_sig_can = TransferMechanism(name="sig_outs_can",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        # SET UP PROJECTIONS FOR SEMANTIC NET

        map_nouns_h1 = MappingProjection(matrix=np.random.rand(8,8),
                                         name="map_nouns_h1",
                                         sender=nouns_in,
                                         receiver=h1)

        map_rels_h2 = MappingProjection(matrix=np.random.rand(3,15),
                                        name="map_relh2",
                                        sender=rels_in,
                                        receiver=h2)

        map_h1_h2 = MappingProjection(matrix=np.random.rand(8,15),
                                      name="map_h1_h2",
                                      sender=h1,
                                      receiver=h2)

        map_h2_I = MappingProjection(matrix=np.random.rand(15,8),
                                     name="map_h2_I",
                                    sender=h2,
                                    receiver=out_sig_I)

        map_h2_is = MappingProjection(matrix=np.random.rand(15,12),
                                      name="map_h2_is",
                                      sender=h2,
                                      receiver=out_sig_is)

        map_h2_has = MappingProjection(matrix=np.random.rand(15,9),
                                       name="map_h2_has",
                                       sender=h2,
                                       receiver=out_sig_has)

        map_h2_can = MappingProjection(matrix=np.random.rand(15,9),
                                       name="map_h2_can",
                                       sender=h2,
                                       receiver=out_sig_can)

        # COMPOSITION FOR SEMANTIC NET
        sem_net = AutodiffComposition(optimizer_type=opt, learning_rate=.001)

        sem_net.add_node(nouns_in)
        sem_net.add_node(rels_in)
        sem_net.add_node(h1)
        sem_net.add_node(h2)
        sem_net.add_node(out_sig_I)
        sem_net.add_node(out_sig_is)
        sem_net.add_node(out_sig_has)
        sem_net.add_node(out_sig_can)

        sem_net.add_projection(sender=nouns_in, projection=map_nouns_h1, receiver=h1)
        sem_net.add_projection(sender=rels_in, projection=map_rels_h2, receiver=h2)
        sem_net.add_projection(sender=h1, projection=map_h1_h2, receiver=h2)
        sem_net.add_projection(sender=h2, projection=map_h2_I, receiver=out_sig_I)
        sem_net.add_projection(sender=h2, projection=map_h2_is, receiver=out_sig_is)
        sem_net.add_projection(sender=h2, projection=map_h2_has, receiver=out_sig_has)
        sem_net.add_projection(sender=h2, projection=map_h2_can, receiver=out_sig_can)

        # INPUTS & OUTPUTS FOR SEMANTIC NET:

        nouns = ['oak', 'pine', 'rose', 'daisy', 'canary', 'robin', 'salmon', 'sunfish']
        relations = ['is', 'has', 'can']
        is_list = ['living', 'living thing', 'plant', 'animal', 'tree', 'flower', 'bird', 'fish', 'big', 'green', 'red',
                   'yellow']
        has_list = ['roots', 'leaves', 'bark', 'branches', 'skin', 'feathers', 'wings', 'gills', 'scales']
        can_list = ['grow', 'move', 'swim', 'fly', 'breathe', 'breathe underwater', 'breathe air', 'walk', 'photosynthesize']

        nouns_input = np.identity(len(nouns))

        rels_input = np.identity(len(relations))

        truth_nouns = np.identity(len(nouns))

        truth_is = np.zeros((len(nouns), len(is_list)))

        truth_is[0, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0]
        truth_is[1, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0]
        truth_is[2, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 1, 0]
        truth_is[3, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 1]
        truth_is[4, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[5, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0]
        truth_is[6, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0]
        truth_is[7, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0]

        truth_has = np.zeros((len(nouns), len(has_list)))

        truth_has[0, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[1, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[2, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[3, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[4, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[5, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[6, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]
        truth_has[7, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]

        truth_can = np.zeros((len(nouns), len(can_list)))

        truth_can[0, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[1, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[2, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[3, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[4, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[5, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[6, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]
        truth_can[7, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]

        # SETTING UP DICTIONARY OF INPUTS/OUTPUTS FOR SEMANTIC NET

        inputs_dict = {}
        inputs_dict[nouns_in] = []
        inputs_dict[rels_in] = []

        targets_dict = {}
        targets_dict[out_sig_I] = []
        targets_dict[out_sig_is] = []
        targets_dict[out_sig_has] = []
        targets_dict[out_sig_can] = []

        for i in range(len(nouns)):
            for j in range(len(relations)):
                inputs_dict[nouns_in].append(nouns_input[i])
                inputs_dict[rels_in].append(rels_input[j])
                targets_dict[out_sig_I].append(truth_nouns[i])
                targets_dict[out_sig_is].append(truth_is[i])
                targets_dict[out_sig_has].append(truth_has[i])
                targets_dict[out_sig_can].append(truth_can[i])

        # TRAIN THE MODEL
        results = benchmark(_single_learn_results,
                            sem_net,
                            inputs={'inputs': inputs_dict,
                                    'targets': targets_dict,
                                    'epochs': eps},
                            execution_mode=autodiff_mode)

        # CHECK CORRECTNESS
        expected = [[[0.13455769, 0.12924714, 0.13288172, 0.1404659 , 0.14305814,
                    0.14801862, 0.13968417, 0.14191982], [0.99991336, 0.99989676, 0.50621126, 0.52845353, 0.25829169,
                    0.26210662, 0.24857125, 0.25041139, 0.4960792 , 0.25689092,
                    0.36730921, 0.24619036], [0.5026992 , 0.51272669, 0.24175637, 0.25034784, 0.2406316 ,
                    0.23613707, 0.25192136, 0.25169522, 0.25615122], [0.99931055, 0.50694136, 0.23105683, 0.24641822, 0.49398116,
                    0.25462159, 0.24433084, 0.24267716, 0.52209116]], [[0.13434049, 0.14280817, 0.13996269, 0.15002735, 0.12975428,
                    0.13143489, 0.14087601, 0.1384075 ], [0.99988782, 0.99987699, 0.50700037, 0.48324852, 0.23605196,
                    0.24623233, 0.25029056, 0.23937042, 0.51119504, 0.2428462 ,
                    0.38335829, 0.24875281], [0.51938115, 0.48541386, 0.25513764, 0.26880953, 0.27165005,
                    0.2768918 , 0.25270764, 0.25394743, 0.2453564 ], [0.99932127, 0.48854068, 0.25824899, 0.24792285, 0.52373987,
                    0.23226666, 0.26078735, 0.25531502, 0.48945013]], [[0.14603474, 0.13076473, 0.14285659, 0.14135699, 0.13505543,
                    0.13515591, 0.14114513, 0.13170369], [0.99990678, 0.99989262, 0.48366278, 0.49653231, 0.24862611,
                    0.23489764, 0.25762623, 0.2557183 , 0.48931577, 0.24042071,
                    0.38135237, 0.2559529 ], [0.48125071, 0.49795935, 0.25190272, 0.23454573, 0.24576829,
                    0.24084843, 0.23960781, 0.24709719, 0.24797553], [0.99927591, 0.50571286, 0.25989537, 0.25337519, 0.49977912,
                    0.2629161 , 0.24592842, 0.25430635, 0.48884817]], [[0.13471359, 0.12938359, 0.12833093, 0.13589542, 0.14362713,
                    0.15158585, 0.14307399, 0.14359844], [0.99991663, 0.99990086, 0.50239329, 0.5363237 , 0.26273315,
                    0.26079003, 0.24686505, 0.2514966 , 0.50092905, 0.26094556,
                    0.36939329, 0.24278735], [0.50633107, 0.51005404, 0.24486936, 0.25305515, 0.2371648 ,
                    0.23443371, 0.24958497, 0.25592672, 0.26077736], [0.99930883, 0.51096502, 0.2273585 , 0.24095976, 0.49420255,
                    0.2541994 , 0.24348927, 0.23870194, 0.52777643]], [[0.13437031, 0.14294567, 0.13462913, 0.14499767, 0.13014469,
                    0.1347783 , 0.14426174, 0.14003815], [0.99989275, 0.99988313, 0.50533481, 0.49070575, 0.24083416,
                    0.24419528, 0.24850006, 0.2388715 , 0.51590916, 0.2477704 ,
                    0.38492997, 0.24480191], [0.52312591, 0.48400766, 0.25827697, 0.27239575, 0.26774108,
                    0.27407821, 0.24982668, 0.25714894, 0.24895088], [0.9993246 , 0.49095133, 0.25329179, 0.2414722 , 0.52311016,
                    0.23199452, 0.25880536, 0.25034698, 0.49623819]], [[0.14624279, 0.13155149, 0.13771617, 0.13640541, 0.1355243 ,
                    0.13836647, 0.14442104, 0.13306906], [0.9999103 , 0.99989714, 0.48256079, 0.50278716, 0.25496791,
                    0.23361944, 0.25567964, 0.25577915, 0.49540433, 0.24632798,
                    0.38186044, 0.25220444], [0.48732786, 0.49713526, 0.25686322, 0.23903771, 0.24179882,
                    0.238869  , 0.2368224 , 0.25066201, 0.25139028], [0.9992754 , 0.50799072, 0.25448626, 0.24711122, 0.49775634,
                    0.26134282, 0.24471328, 0.24960873, 0.49669515]], [[0.13743409, 0.12655781, 0.13687131, 0.14444673, 0.14586922,
                    0.14982043, 0.14115407, 0.14425256], [0.999905  , 0.99988946, 0.5111958 , 0.52953632, 0.25829416,
                    0.25816645, 0.24772719, 0.2493831 , 0.49172307, 0.25647241,
                    0.36490702, 0.24579828], [0.49409604, 0.51876333, 0.2438246 , 0.25106641, 0.24575334,
                    0.23758284, 0.25036778, 0.24501611, 0.26348273], [0.99927212, 0.49900832, 0.23785606, 0.24568878, 0.49229404,
                    0.25799455, 0.2380617 , 0.24483316, 0.51242196]], [[0.13690143, 0.14031709, 0.14406563, 0.15349375, 0.13257716,
                    0.13310877, 0.14252974, 0.14056437], [0.99987794, 0.99986931, 0.5140066 , 0.48347409, 0.23776129,
                    0.24271912, 0.2491253 , 0.23773346, 0.50725307, 0.24412832,
                    0.37955962, 0.24840394], [0.5134345 , 0.49359454, 0.25820627, 0.27104423, 0.27620957,
                    0.27795555, 0.25014143, 0.24655665, 0.25137209], [0.99928701, 0.47923949, 0.26354622, 0.24637029, 0.51954436,
                    0.23441734, 0.25357643, 0.25643598, 0.48217473]], [[0.14885265, 0.12888511, 0.14747413, 0.14492694, 0.13757074,
                    0.13653999, 0.14230078, 0.1333605 ], [0.99989882, 0.99988608, 0.49175208, 0.49464422, 0.25033443,
                    0.2315963 , 0.25565926, 0.25286519, 0.48521686, 0.24180012,
                    0.37838253, 0.25482498], [0.47632935, 0.50754291, 0.25540722, 0.23699234, 0.24989352,
                    0.24130117, 0.236814  , 0.23865942, 0.25337658], [0.99924137, 0.49439904, 0.26500386, 0.25133608, 0.49471845,
                    0.26444086, 0.23836089, 0.2550785 , 0.48265071]], [[0.13796428, 0.13122133, 0.13665485, 0.14238789, 0.14646363,
                    0.14953325, 0.14001922, 0.14631467], [0.99990658, 0.99988831, 0.51044886, 0.52877409, 0.26162535,
                    0.26113322, 0.24464028, 0.25460667, 0.50142119, 0.26075076,
                    0.36249296, 0.2469639 ], [0.49888045, 0.520344  , 0.24213434, 0.25386651, 0.23856703,
                    0.23283943, 0.25382764, 0.25348876, 0.25930602], [0.9992759 , 0.50552098, 0.23581279, 0.25056642, 0.49046455,
                    0.26042018, 0.24011749, 0.24306269, 0.52240185]], [[0.13750907, 0.14541291, 0.1443253 , 0.15189612, 0.13295338,
                    0.13279175, 0.14102944, 0.14247126], [0.99987945, 0.99986745, 0.51385772, 0.48177978, 0.23992435,
                    0.24651741, 0.24564878, 0.2423684 , 0.51605876, 0.24748102,
                    0.37804356, 0.24958671], [0.51825942, 0.49563748, 0.25600593, 0.27323944, 0.26870607,
                    0.27273737, 0.25405544, 0.25468467, 0.24732677], [0.99928937, 0.4851184 , 0.26178837, 0.25136875, 0.51760249,
                    0.23648163, 0.25573403, 0.25496827, 0.49227648]], [[0.14930882, 0.13316748, 0.14767505, 0.14350067, 0.13804021,
                    0.13603459, 0.14084128, 0.13513649], [0.99990066, 0.99988511, 0.49193637, 0.49266043, 0.25237073,
                    0.23599757, 0.25202475, 0.25740929, 0.49266556, 0.24476454,
                    0.37618961, 0.25697362], [0.48210255, 0.51007883, 0.25286726, 0.2385758 , 0.24213314,
                    0.23616873, 0.23972905, 0.24636349, 0.24902741], [0.99924664, 0.4998954 , 0.2624679 , 0.25622887, 0.49190217,
                    0.26610712, 0.24000543, 0.25322914, 0.49359824]], [[0.13981779, 0.13517728, 0.13714904, 0.14312451, 0.14264037,
                    0.1471793 , 0.13872158, 0.14604209], [0.99990026, 0.99988117, 0.50978024, 0.51828715, 0.26366737,
                    0.26281186, 0.24835743, 0.25797436, 0.50770859, 0.26547003,
                    0.36335027, 0.24615446], [0.50366793, 0.51568697, 0.2470342 , 0.2533224 , 0.23632146,
                    0.23162517, 0.25117307, 0.25934756, 0.256024  ], [0.99924436, 0.50891146, 0.23908499, 0.24926875, 0.48979495,
                    0.26367751, 0.24330774, 0.24296292, 0.523776  ]], [[0.13961147, 0.14975401, 0.14472653, 0.15355384, 0.12925548,
                    0.13046081, 0.13938395, 0.14208934], [0.99987039, 0.99985783, 0.51256916, 0.47090975, 0.24108525,
                    0.24849742, 0.24950535, 0.24553023, 0.522052  , 0.25148739,
                    0.37924791, 0.24931395], [0.52265254, 0.49040645, 0.26115407, 0.27231193, 0.26698365,
                    0.27201252, 0.252447  , 0.26033848, 0.24417141], [0.99925421, 0.48866289, 0.26591473, 0.25072477, 0.51818878,
                    0.23925897, 0.26008435, 0.25565447, 0.49313837]], [[0.15131984, 0.13702244, 0.14796541, 0.14499107, 0.13516231,
                    0.13379344, 0.13967548, 0.13504357], [0.99989266, 0.99987644, 0.49069576, 0.48284025, 0.25373136,
                    0.23780959, 0.25705167, 0.26031535, 0.49734223, 0.24868116,
                    0.37609291, 0.25826094], [0.48615805, 0.50439585, 0.25764469, 0.23792093, 0.24102865,
                    0.23622676, 0.23835519, 0.25200694, 0.24543424], [0.99920779, 0.50378559, 0.265772  , 0.25635798, 0.49252096,
                    0.26892098, 0.24458074, 0.25443214, 0.494262  ]], [[0.13850469, 0.13222947, 0.13670674, 0.14224276, 0.14481382,
                    0.14983366, 0.14173916, 0.14323358], [0.99990126, 0.99988249, 0.51196132, 0.52177217, 0.2678038 ,
                    0.26687693, 0.24660127, 0.25415374, 0.50244534, 0.26499807,
                    0.36556967, 0.24445311], [0.51065084, 0.52060372, 0.25297718, 0.25785385, 0.24041214,
                    0.23595817, 0.24912937, 0.25285752, 0.26092579], [0.9992282 , 0.49866166, 0.23316642, 0.2414636 , 0.48397547,
                    0.25549083, 0.24296723, 0.24121413, 0.52743639]], [[0.13795671, 0.14608656, 0.14389815, 0.15250863, 0.13137121,
                    0.13292996, 0.14235728, 0.13916877], [0.99987232, 0.99986024, 0.514095  , 0.47542769, 0.24452563,
                    0.25126078, 0.248764  , 0.241217  , 0.51579474, 0.25047855,
                    0.38175049, 0.24768651], [0.52795513, 0.49433932, 0.26639626, 0.2765692 , 0.27249335,
                    0.27735353, 0.25105699, 0.25316897, 0.24833112], [0.99924257, 0.47881092, 0.25918246, 0.24338929, 0.51375801,
                    0.23187301, 0.26019419, 0.25420934, 0.49555353]], [[0.14984495, 0.13387861, 0.14722304, 0.14355772, 0.13724767,
                    0.13665399, 0.14239016, 0.1320989 ], [0.99989382, 0.9998781 , 0.49129168, 0.48786117, 0.25736056,
                    0.24076914, 0.2571532 , 0.25661644, 0.49134559, 0.24762313,
                    0.37894905, 0.25678541], [0.49115594, 0.50720606, 0.26343239, 0.24118617, 0.24666918,
                    0.24228136, 0.23828715, 0.24556003, 0.24986447], [0.9991909 , 0.49609669, 0.25952719, 0.2503092 , 0.48831756,
                    0.26050785, 0.24624062, 0.25460693, 0.49571571]], [[0.13798156, 0.13536763, 0.13556549, 0.14026479, 0.14613495,
                    0.1517147 , 0.14150378, 0.14582109], [0.99990224, 0.99988499, 0.51321887, 0.5234933 , 0.26596025,
                    0.26627399, 0.24845062, 0.25224455, 0.50060705, 0.26840742,
                    0.36382243, 0.24724158], [0.50857345, 0.52097167, 0.24807872, 0.25421915, 0.23896201,
                    0.23521938, 0.25167447, 0.25399296, 0.25680194], [0.99924088, 0.50538628, 0.23168769, 0.24466784, 0.48374493,
                    0.25836643, 0.24395082, 0.24271583, 0.52908765]], [[0.13738911, 0.14964169, 0.14240066, 0.14996998, 0.13258046,
                    0.13529629, 0.14257416, 0.14196438], [0.99987348, 0.99986355, 0.51547364, 0.47846669, 0.24280831,
                    0.25045232, 0.25106236, 0.23941227, 0.51388317, 0.25405565,
                    0.38069812, 0.25010755], [0.52466269, 0.49422476, 0.26072088, 0.27243928, 0.27097311,
                    0.27644019, 0.25390417, 0.25487688, 0.24468432], [0.99925402, 0.48653142, 0.25783738, 0.24674309, 0.51370588,
                    0.23544836, 0.26119428, 0.25606803, 0.49686813]], [[0.14925705, 0.13686349, 0.14561612, 0.14139895, 0.13823102,
                    0.13876809, 0.14286744, 0.13452957], [0.99989461, 0.99988054, 0.49152081, 0.49115061, 0.25525704,
                    0.23961845, 0.25940244, 0.25531285, 0.48988801, 0.25049358,
                    0.37898184, 0.25880442], [0.48716519, 0.50598831, 0.25786953, 0.23748522, 0.24522105,
                    0.24133955, 0.24092032, 0.247859  , 0.24660646], [0.99920254, 0.50397091, 0.25882576, 0.25383231, 0.48962604,
                    0.26477928, 0.24738806, 0.25637011, 0.49634854]], [[0.13466774, 0.13597299, 0.13044001, 0.13621353, 0.14063205,
                    0.14927389, 0.13741799, 0.14182897], [0.99991209, 0.99989375, 0.50664715, 0.52647149, 0.26631381,
                    0.26772046, 0.25065527, 0.25828284, 0.51113304, 0.26876027,
                    0.36519173, 0.24445213], [0.50981376, 0.51216926, 0.24469315, 0.25069097, 0.23575157,
                    0.23351051, 0.25738456, 0.26157098, 0.25154516], [0.9992887 , 0.51923276, 0.2286251 , 0.24647828, 0.48904116,
                    0.25948029, 0.25191475, 0.24172423, 0.53116312]], [[0.13456692, 0.15039528, 0.13737152, 0.14609287, 0.12736956,
                    0.13283867, 0.13853673, 0.13816526], [0.99988542, 0.99987272, 0.50746352, 0.48137383, 0.24278802,
                    0.25182679, 0.25260023, 0.24645783, 0.52561721, 0.25394353,
                    0.38280036, 0.24682469], [0.52522447, 0.48470641, 0.25775314, 0.26859876, 0.26712899,
                    0.27417417, 0.25948482, 0.26348026, 0.24082551], [0.9992959 , 0.5009347 , 0.25614941, 0.24856744, 0.51984942,
                    0.23715273, 0.2693644 , 0.2551465 , 0.49819368]], [[0.14562679, 0.13745423, 0.13996431, 0.13722137, 0.13278566,
                    0.13632886, 0.13899164, 0.1314741 ], [0.99990502, 0.99988943, 0.48410876, 0.49522954, 0.2553264 ,
                    0.24084323, 0.26058312, 0.26304712, 0.50254737, 0.2507632 ,
                    0.3802493 , 0.25486957], [0.48768232, 0.49615974, 0.25448472, 0.23400199, 0.24116754,
                    0.23886036, 0.24575353, 0.25715595, 0.24334699], [0.99925183, 0.51889063, 0.25712839, 0.25460613, 0.49597306,
                    0.26739429, 0.25464059, 0.25453138, 0.49761396]]]


        if pytest.helpers.llvm_current_fp_precision() == 'fp32' and autodiff_mode != pnl.ExecutionMode.PyTorch:
            accuracy_args = {"atol": 1e-8, "rtol": 1e-6}
        else:
            accuracy_args = {}

        # The results are in a ragged array, we need to check components explictly
        for res, exp in zip(results, expected):
            for r, e in zip(res, exp):
                np.testing.assert_allclose(r, e, **accuracy_args)


    def test_pytorch_equivalence_with_autodiff_composition(self, autodiff_mode):
        iSs = np.array(
            [np.array([0.47360805, 0.8009108, 0.5204775, 0.53737324, 0.7586156,
                    0.1059076, 0.9025985, 0.44994998, 0.61306345, 0.75068617,
                    0.60783064, 0.32504722, 0.58185035, 0.4143686, 0.4746975]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                    0.1059076, 0.21655035, 0.13521817, 0.324141, 0.65314,
                    0.17090958, 0.35815218, 0.03842543, 0.63427407, 0.95894927]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                    0.1059076, 0.9025985, 0.44994998, 0.61306345, 0.65314,
                    0.17090958, 0.35815218, 0.58185035, 0.4143686, 0.4746975]),
             np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                    0.73691815, 0.14967486, 0.22232139, 0.38648897, 0.75068617,
                    0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
             np.array([0.47360805, 0.8009108, 0.5204775, 0.47360042, 0.18633235,
                    0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                    0.09928035, 0.96980906, 0.03842543, 0.63427407, 0.95894927]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                    0.1059076, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                    0.09928035, 0.96980906, 0.03842543, 0.63427407, 0.95894927]),
             np.array([0.47360805, 0.8009108, 0.5204775, 0.47360042, 0.18633235,
                    0.73691815, 0.14967486, 0.22232139, 0.38648897, 0.65314,
                    0.17090958, 0.35815218, 0.58185035, 0.4143686, 0.4746975]),
             np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                    0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.75068617,
                    0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
             np.array([0.47360805, 0.8009108, 0.5204775, 0.53737324, 0.7586156,
                    0.1059076, 0.21655035, 0.13521817, 0.324141, 0.75068617,
                    0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                    0.1059076, 0.14967486, 0.22232139, 0.38648897, 0.9023486,
                    0.09928035, 0.96980906, 0.03842543, 0.63427407, 0.95894927]),
             np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                    0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                    0.09928035, 0.96980906, 0.6527903, 0.6350589, 0.9952996]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.47360042, 0.18633235,
                    0.73691815, 0.14967486, 0.22232139, 0.38648897, 0.75068617,
                    0.60783064, 0.32504722, 0.03842543, 0.63427407, 0.95894927]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                    0.1059076, 0.14967486, 0.22232139, 0.38648897, 0.65314,
                    0.17090958, 0.35815218, 0.03842543, 0.63427407, 0.95894927]),
             np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                    0.73691815, 0.21655035, 0.13521817, 0.324141, 0.75068617,
                    0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
             np.array([0.33739617, 0.6481719, 0.36824155, 0.47360042, 0.18633235,
                    0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                    0.09928035, 0.96980906, 0.6527903, 0.6350589, 0.9952996])]
        )

        cSs = np.array(
            [np.array([0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 1., 0., 0.]),
             np.array([0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 1.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 1., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 1., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 1.]),
             np.array([0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                    0., 0., 0., 0., 0., 0., 0., 1.])]
        )

        oSs = np.array(
            [np.array([0., 1., -0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([1., 0., -0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., -0., 0., 1., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., -0., 0.]),
             np.array([0., 0., 0., 0., 1., -0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., -0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., -0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., -0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 1., -0., 0., 0., 0., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., -0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.]),
             np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.])]
        )

        nf = 3
        nd = 5
        nh = 200

        D_i = nf * nd
        D_c = nd ** 2
        D_h = nh
        D_o = nf * nd

        np.random.seed(0)

        wih = np.random.rand(D_i, D_h) * 0.02 - 0.01
        wch = np.random.rand(D_c, D_h) * 0.02 - 0.01
        wco = np.random.rand(D_c, D_o) * 0.02 - 0.01
        who = np.random.rand(D_h, D_o) * 0.02 - 0.01

        patience = 10
        min_delt = 0.00001
        learning_rate = 100

        il = TransferMechanism(input_shapes=D_i, name='input')
        cl = TransferMechanism(input_shapes=D_c, name='task')
        hl = TransferMechanism(
            input_shapes=D_h, name='hidden',
                               function=Logistic(bias=-2))
        ol = TransferMechanism(
            input_shapes=D_o, name='output',
                               function=Logistic(bias=-2))

        input_set = {
            'inputs': {
                il: iSs,
                cl: cSs
            },
            'targets': {
                ol: oSs
            }
        }

        pih = MappingProjection(matrix=wih)
        pch = MappingProjection(matrix=wch)
        pco = MappingProjection(matrix=wco)
        pho = MappingProjection(matrix=who)

        mnet = AutodiffComposition(learning_rate=learning_rate)

        mnet.add_node(il)
        mnet.add_node(cl)
        mnet.add_node(hl)
        mnet.add_node(ol)
        mnet.add_projection(projection=pih, sender=il, receiver=hl)
        mnet.add_projection(projection=pch, sender=cl, receiver=hl)
        mnet.add_projection(projection=pco, sender=cl, receiver=ol)
        mnet.add_projection(projection=pho, sender=hl, receiver=ol)

        mnet.learn(
            inputs=input_set,
            minibatch_size=1,
            patience=patience,
            min_delta=min_delt,
            execution_mode=autodiff_mode
        )
        mnet.run(
            inputs=input_set['inputs'],
            execution_mode=autodiff_mode
        )
        output = np.array(mnet.parameters.results.get(mnet)[-15:]).reshape(225)

        comparator = np.array([0.02288846, 0.11646781, 0.03473711, 0.0348004, 0.01679579,
                             0.04851733, 0.05857743, 0.04819957, 0.03004438, 0.05113508,
                             0.06849843, 0.0442623, 0.00967315, 0.06998125, 0.03482444,
                             0.05856816, 0.00724313, 0.03676571, 0.03668758, 0.01761947,
                             0.0516829, 0.06260267, 0.05160782, 0.03140498, 0.05462971,
                             0.07360401, 0.04687923, 0.00993319, 0.07662302, 0.03687142,
                             0.0056837, 0.03411045, 0.03615285, 0.03606166, 0.01774354,
                             0.04700402, 0.09696857, 0.06843472, 0.06108671, 0.0485631,
                             0.07194324, 0.04485926, 0.00526768, 0.07442083, 0.0364541,
                             0.02819926, 0.03804169, 0.04091214, 0.04091113, 0.04246229,
                             0.05583883, 0.06643675, 0.05630667, 0.01540373, 0.05948422,
                             0.07721549, 0.05081813, 0.01205326, 0.07998289, 0.04084186,
                             0.02859247, 0.03794089, 0.04111452, 0.04139213, 0.01222424,
                             0.05677404, 0.06736114, 0.05614553, 0.03573626, 0.05983103,
                             0.07867571, 0.09971621, 0.01203033, 0.08107789, 0.04110497,
                             0.02694072, 0.03592752, 0.03878366, 0.03895513, 0.01852774,
                             0.05097689, 0.05753834, 0.05090328, 0.03405996, 0.05293719,
                             0.07037981, 0.03474316, 0.02861534, 0.12504038, 0.0387827,
                             0.02467716, 0.03373265, 0.03676382, 0.03677551, 0.00758558,
                             0.089832, 0.06330426, 0.0514472, 0.03120581, 0.05535174,
                             0.07494839, 0.04169744, 0.00698747, 0.0771042, 0.03659954,
                             0.03008443, 0.0393799, 0.0423592, 0.04237004, 0.00965198,
                             0.09863199, 0.06813933, 0.05675321, 0.03668943, 0.0606036,
                             0.07898065, 0.04662618, 0.00954765, 0.08093391, 0.04218842,
                             0.02701085, 0.03660227, 0.04058368, 0.04012464, 0.02030738,
                             0.047633, 0.06693405, 0.055821, 0.03456592, 0.10166267,
                             0.07870758, 0.04935871, 0.01065449, 0.08012213, 0.04036544,
                             0.02576563, 0.03553382, 0.03920509, 0.03914452, 0.01907667,
                             0.05106766, 0.06555857, 0.05434728, 0.03335726, 0.05074808,
                             0.07715102, 0.04839309, 0.02494798, 0.08001304, 0.03921895,
                             0.00686952, 0.03941704, 0.04128484, 0.04117602, 0.02217508,
                             0.05152296, 0.10361618, 0.07488737, 0.0707186, 0.05289282,
                             0.07557573, 0.04978292, 0.00705783, 0.07787788, 0.04164007,
                             0.00574239, 0.03437231, 0.03641445, 0.03631848, 0.01795791,
                             0.04723996, 0.09732232, 0.06876138, 0.06156679, 0.04878423,
                             0.07214104, 0.04511085, 0.00535038, 0.07459818, 0.0367153,
                             0.02415251, 0.03298647, 0.03586635, 0.0360273, 0.01624523,
                             0.04829838, 0.05523439, 0.04821285, 0.03115052, 0.05034625,
                             0.06836408, 0.03264844, 0.0241706, 0.12190507, 0.03585727,
                             0.02897192, 0.03925683, 0.04250414, 0.04253885, 0.02175426,
                             0.05683923, 0.06547528, 0.05705267, 0.03742978, 0.05951711,
                             0.12675475, 0.05216411, 0.00181494, 0.08218002, 0.04234364,
                             0.02789848, 0.036924, 0.03976586, 0.03993866, 0.01932489,
                             0.05186586, 0.05829845, 0.05179337, 0.03504668, 0.05379566,
                             0.07103772, 0.03544133, 0.03019486, 0.12605846, 0.03976812])

        np.testing.assert_allclose(output,comparator, atol=1e-6)

    def test_pytorch_equivalence_with_autodiff_forward_disabled_on_proj(self):
        iSs = np.array(
                [np.array([0.47360805, 0.8009108, 0.5204775, 0.53737324, 0.7586156,
                           0.1059076, 0.9025985, 0.44994998, 0.61306345, 0.75068617,
                           0.60783064, 0.32504722, 0.58185035, 0.4143686, 0.4746975]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                           0.1059076, 0.21655035, 0.13521817, 0.324141, 0.65314,
                           0.17090958, 0.35815218, 0.03842543, 0.63427407, 0.95894927]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                           0.1059076, 0.9025985, 0.44994998, 0.61306345, 0.65314,
                           0.17090958, 0.35815218, 0.58185035, 0.4143686, 0.4746975]),
                 np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                           0.73691815, 0.14967486, 0.22232139, 0.38648897, 0.75068617,
                           0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
                 np.array([0.47360805, 0.8009108, 0.5204775, 0.47360042, 0.18633235,
                           0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                           0.09928035, 0.96980906, 0.03842543, 0.63427407, 0.95894927]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                           0.1059076, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                           0.09928035, 0.96980906, 0.03842543, 0.63427407, 0.95894927]),
                 np.array([0.47360805, 0.8009108, 0.5204775, 0.47360042, 0.18633235,
                           0.73691815, 0.14967486, 0.22232139, 0.38648897, 0.65314,
                           0.17090958, 0.35815218, 0.58185035, 0.4143686, 0.4746975]),
                 np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                           0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.75068617,
                           0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
                 np.array([0.47360805, 0.8009108, 0.5204775, 0.53737324, 0.7586156,
                           0.1059076, 0.21655035, 0.13521817, 0.324141, 0.75068617,
                           0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                           0.1059076, 0.14967486, 0.22232139, 0.38648897, 0.9023486,
                           0.09928035, 0.96980906, 0.03842543, 0.63427407, 0.95894927]),
                 np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                           0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                           0.09928035, 0.96980906, 0.6527903, 0.6350589, 0.9952996]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.47360042, 0.18633235,
                           0.73691815, 0.14967486, 0.22232139, 0.38648897, 0.75068617,
                           0.60783064, 0.32504722, 0.03842543, 0.63427407, 0.95894927]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.53737324, 0.7586156,
                           0.1059076, 0.14967486, 0.22232139, 0.38648897, 0.65314,
                           0.17090958, 0.35815218, 0.03842543, 0.63427407, 0.95894927]),
                 np.array([0.95715517, 0.14035077, 0.87008727, 0.47360042, 0.18633235,
                           0.73691815, 0.21655035, 0.13521817, 0.324141, 0.75068617,
                           0.60783064, 0.32504722, 0.6527903, 0.6350589, 0.9952996]),
                 np.array([0.33739617, 0.6481719, 0.36824155, 0.47360042, 0.18633235,
                           0.73691815, 0.9025985, 0.44994998, 0.61306345, 0.9023486,
                           0.09928035, 0.96980906, 0.6527903, 0.6350589, 0.9952996])]
        )

        cSs = np.array(
                [np.array([0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 1., 0., 0.]),
                 np.array([0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 1.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 1., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 1., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 1.]),
                 np.array([0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
                           0., 0., 0., 0., 0., 0., 0., 1.])]
        )

        oSs = np.array(
                [np.array([0., 1., -0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([1., 0., -0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., -0., 0., 1., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., -0., 0.]),
                 np.array([0., 0., 0., 0., 1., -0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., -0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., -0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., -0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 1., -0., 0., 0., 0., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., -0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., 0., 0., 0.]),
                 np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 0.])]
        )

        nf = 3
        nd = 5
        nh = 200

        D_i = nf * nd
        D_c = nd ** 2
        D_h = nh
        D_o = nf * nd

        np.random.seed(0)

        wih = np.random.rand(D_i, D_h) * 0.02 - 0.01
        wch = np.random.rand(D_c, D_h) * 0.02 - 0.01
        wco = np.random.rand(D_c, D_o) * 0.02 - 0.01
        who = np.random.rand(D_h, D_o) * 0.02 - 0.01

        patience = 10
        min_delt = 0.00001
        learning_rate = 100

        il = TransferMechanism(input_shapes=D_i, name='input')
        cl = TransferMechanism(input_shapes=D_c, name='task')
        hl = TransferMechanism(
            input_shapes=D_h, name='hidden',
                               function=Logistic(bias=-2))
        ol = TransferMechanism(
            input_shapes=D_o, name='output',
                               function=Logistic(bias=-2))

        input_set = {
            'inputs': {
                il: iSs,
                cl: cSs
            },
            'targets': {
                ol: oSs
            }
        }

        pih = MappingProjection(matrix=wih, name='pih')
        pch = MappingProjection(matrix=wch, name='pch')
        pco = MappingProjection(matrix=wco, name='pco')
        pho = MappingProjection(matrix=who, name='pho', learnable=False)

        mnet = AutodiffComposition(learning_rate=learning_rate)

        mnet.add_node(il)
        mnet.add_node(cl)
        mnet.add_node(hl)
        mnet.add_node(ol)
        mnet.add_projection(projection=pih, sender=il, receiver=hl)
        mnet.add_projection(projection=pch, sender=cl, receiver=hl)
        mnet.add_projection(projection=pco, sender=cl, receiver=ol)
        mnet.add_projection(projection=pho, sender=hl, receiver=ol)


        result = mnet.learn(
            inputs=input_set,
            minibatch_size=1,
            patience=patience,
            min_delta=min_delt,
            execution_mode=pnl.ExecutionMode.PyTorch,
        )

        mnet.run(
            inputs=input_set['inputs']
        )

        output = np.array(mnet.parameters.results.get(mnet)[-15:]).reshape(225)

        comparator = np.array([0.10284232, 0.31514028, 0.10299414, 0.10164745, 0.10363132,
                               0.10164711, 0.10305342, 0.10162935, 0.10363974, 0.10175142,
                               0.10256631, 0.10194203, 0.10386363, 0.10445295, 0.10228054,
                               0.31140432, 0.10257346, 0.10279541, 0.1015088, 0.10408029,
                               0.10167408, 0.10260046, 0.10208146, 0.10258093, 0.10188455,
                               0.10239721, 0.10162553, 0.10376681, 0.10523887, 0.10231788,
                               0.08327345, 0.08337342, 0.0835557, 0.0828431, 0.08364569,
                               0.08285296, 0.21721269, 0.15223454, 0.12355195, 0.08328209,
                               0.08321026, 0.08318614, 0.08401372, 0.08443127, 0.08355132,
                               0.10225081, 0.10250866, 0.1032809, 0.10216374, 0.3212671,
                               0.10171002, 0.10223842, 0.10279202, 0.10348979, 0.102771,
                               0.10200755, 0.10137874, 0.10408875, 0.10449553, 0.10241774,
                               0.10293344, 0.10201894, 0.10300561, 0.10239243, 0.10291971,
                               0.10242151, 0.10280451, 0.10199619, 0.10344362, 0.10265052,
                               0.1030072, 0.31077573, 0.10299222, 0.10510338, 0.10226066,
                               0.08338644, 0.08334018, 0.08376527, 0.08334996, 0.08397464,
                               0.08293792, 0.08313457, 0.08310839, 0.08409815, 0.08289795,
                               0.08348748, 0.08323742, 0.35242194, 0.22024544, 0.08337309,
                               0.09164643, 0.09135997, 0.09195332, 0.09117354, 0.15678808,
                               0.25366357, 0.09192788, 0.09090009, 0.09173747, 0.09161069,
                               0.09198699, 0.09058125, 0.09191367, 0.09321109, 0.09121469,
                               0.09163069, 0.09134816, 0.09194396, 0.09114014, 0.15678652,
                               0.2536617, 0.09192093, 0.09089337, 0.09171399, 0.09160125,
                               0.09198645, 0.09058312, 0.09191372, 0.09321296, 0.09118975,
                               0.10222919, 0.1017347, 0.10354281, 0.10158797, 0.1038858,
                               0.10181702, 0.10269418, 0.10235615, 0.10275149, 0.31305784,
                               0.1030191, 0.10225646, 0.10283817, 0.10411466, 0.10244074,
                               0.10203665, 0.10201294, 0.10314981, 0.10192659, 0.10328009,
                               0.10265024, 0.1021864, 0.10181551, 0.1026119, 0.10268809,
                               0.10219657, 0.10172481, 0.32032955, 0.104648, 0.10248389,
                               0.08325538, 0.08334755, 0.08355319, 0.08281158, 0.08365688,
                               0.08285309, 0.21719442, 0.15221967, 0.12351983, 0.08326486,
                               0.08321615, 0.08318119, 0.08400558, 0.0844217, 0.08352901,
                               0.08326998, 0.08336743, 0.08356269, 0.08283862, 0.08365061,
                               0.08286179, 0.21723635, 0.15221801, 0.12355236, 0.08327687,
                               0.08322325, 0.08318282, 0.08401041, 0.08442231, 0.0835505,
                               0.0833958, 0.08335006, 0.08376891, 0.08336972, 0.08397432,
                               0.08294199, 0.08314709, 0.08311359, 0.0841146, 0.08291036,
                               0.08349533, 0.08323479, 0.35241473, 0.22023965, 0.08338647,
                               0.10243648, 0.10270733, 0.10287204, 0.10181676, 0.10309494,
                               0.10208003, 0.10258352, 0.10279328, 0.10355093, 0.10241994,
                               0.31674582, 0.10140157, 0.10286999, 0.10426361, 0.1018871,
                               0.08337424, 0.08333415, 0.08376191, 0.08333433, 0.08398008,
                               0.08293781, 0.08313539, 0.08310112, 0.08409653, 0.08289441,
                               0.08348761, 0.08323367, 0.35237628, 0.22024095, 0.08336799])

        np.testing.assert_allclose(output, comparator, rtol=1e-5, atol=1e-8)


@pytest.mark.pytorch
@pytest.mark.acidenticalness
class TestTrainingIdenticalness():

    @pytest.mark.parametrize(
        'eps, opt', [
            # (1, 'sgd'),
            (10, 'sgd'),
            # (40, 'sgd')
        ]
    )
    def test_semantic_net_training_identicalness(self, eps, opt):
        # SET UP MECHANISMS FOR SEMANTIC NET:

        nouns_in = TransferMechanism(name="nouns_input",
                                     default_variable=np.zeros(8))

        rels_in = TransferMechanism(name="rels_input",
                                    default_variable=np.zeros(3))

        h1 = TransferMechanism(name="hidden_nouns",
                               default_variable=np.zeros(8),
                               function=Logistic())

        h2 = TransferMechanism(name="hidden_mixed",
                               default_variable=np.zeros(15),
                               function=Logistic())

        out_sig_I = TransferMechanism(name="sig_outs_I",
                                      default_variable=np.zeros(8),
                                      function=Logistic())

        out_sig_is = TransferMechanism(name="sig_outs_is",
                                       default_variable=np.zeros(12),
                                       function=Logistic())

        out_sig_has = TransferMechanism(name="sig_outs_has",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        out_sig_can = TransferMechanism(name="sig_outs_can",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        # SET UP MECHANISMS FOR Composition

        nouns_in_comp = TransferMechanism(name="nouns_input_comp",
                                         default_variable=np.zeros(8))

        rels_in_comp = TransferMechanism(name="rels_input_comp",
                                        default_variable=np.zeros(3))

        h1_comp = TransferMechanism(name="hidden_nouns_comp",
                                   default_variable=np.zeros(8),
                                   function=Logistic())

        h2_comp = TransferMechanism(name="hidden_mixed_comp",
                                   default_variable=np.zeros(15),
                                   function=Logistic())

        out_sig_I_comp = TransferMechanism(name="sig_outs_I_comp",
                                          default_variable=np.zeros(8),
                                          function=Logistic())

        out_sig_is_comp = TransferMechanism(name="sig_outs_is_comp",
                                           default_variable=np.zeros(12),
                                           function=Logistic())

        out_sig_has_comp = TransferMechanism(name="sig_outs_has_comp",
                                            default_variable=np.zeros(9),
                                            function=Logistic())

        out_sig_can_comp = TransferMechanism(name="sig_outs_can_comp",
                                            default_variable=np.zeros(9),
                                            function=Logistic())

        # SET UP PROJECTIONS FOR SEMANTIC NET

        map_nouns_h1 = MappingProjection(matrix=np.random.rand(8,8),
                                 name="map_nouns_h1",
                                 sender=nouns_in,
                                 receiver=h1)

        map_rels_h2 = MappingProjection(matrix=np.random.rand(3,15),
                                    name="map_relh2",
                                    sender=rels_in,
                                    receiver=h2)

        map_h1_h2 = MappingProjection(matrix=np.random.rand(8,15),
                                    name="map_h1_h2",
                                    sender=h1,
                                    receiver=h2)

        map_h2_I = MappingProjection(matrix=np.random.rand(15,8),
                                    name="map_h2_I",
                                    sender=h2,
                                    receiver=out_sig_I)

        map_h2_is = MappingProjection(matrix=np.random.rand(15,12),
                                    name="map_h2_is",
                                    sender=h2,
                                    receiver=out_sig_is)

        map_h2_has = MappingProjection(matrix=np.random.rand(15,9),
                                    name="map_h2_has",
                                    sender=h2,
                                    receiver=out_sig_has)

        map_h2_can = MappingProjection(matrix=np.random.rand(15,9),
                                    name="map_h2_can",
                                    sender=h2,
                                    receiver=out_sig_can)

        # SET UP PROJECTIONS FOR COMPOSITION

        map_nouns_h1_comp = MappingProjection(matrix=map_nouns_h1.matrix.base.copy(),
                                             name="map_nouns_h1_comp",
                                             sender=nouns_in_comp,
                                             receiver=h1_comp)

        map_rels_h2_comp = MappingProjection(matrix=map_rels_h2.matrix.base.copy(),
                                        name="map_relh2_comp",
                                        sender=rels_in_comp,
                                        receiver=h2_comp)

        map_h1_h2_comp = MappingProjection(matrix=map_h1_h2.matrix.base.copy(),
                                          name="map_h1_h2_comp",
                                          sender=h1_comp,
                                          receiver=h2_comp)

        map_h2_I_comp = MappingProjection(matrix=map_h2_I.matrix.base.copy(),
                                         name="map_h2_I_comp",
                                         sender=h2_comp,
                                         receiver=out_sig_I_comp)

        map_h2_is_comp = MappingProjection(matrix=map_h2_is.matrix.base.copy(),
                                          name="map_h2_is_comp",
                                          sender=h2_comp,
                                          receiver=out_sig_is_comp)

        map_h2_has_comp = MappingProjection(matrix=map_h2_has.matrix.base.copy(),
                                           name="map_h2_has_comp",
                                           sender=h2_comp,
                                           receiver=out_sig_has_comp)

        map_h2_can_comp = MappingProjection(matrix=map_h2_can.matrix.base.copy(),
                                           name="map_h2_can_comp",
                                           sender=h2_comp,
                                           receiver=out_sig_can_comp)

        # SET UP AUTODIFFCOMPOSITION FOR SEMANTIC NET
        sem_net_autodiff = AutodiffComposition(learning_rate=0.5,
                                      optimizer_type=opt,
                                      )

        sem_net_autodiff.add_node(nouns_in)
        sem_net_autodiff.add_node(rels_in)
        sem_net_autodiff.add_node(h1)
        sem_net_autodiff.add_node(h2)
        sem_net_autodiff.add_node(out_sig_I)
        sem_net_autodiff.add_node(out_sig_is)
        sem_net_autodiff.add_node(out_sig_has)
        sem_net_autodiff.add_node(out_sig_can)

        sem_net_autodiff.add_projection(sender=nouns_in, projection=map_nouns_h1, receiver=h1)
        sem_net_autodiff.add_projection(sender=rels_in, projection=map_rels_h2, receiver=h2)
        sem_net_autodiff.add_projection(sender=h1, projection=map_h1_h2, receiver=h2)
        sem_net_autodiff.add_projection(sender=h2, projection=map_h2_I, receiver=out_sig_I)
        sem_net_autodiff.add_projection(sender=h2, projection=map_h2_is, receiver=out_sig_is)
        sem_net_autodiff.add_projection(sender=h2, projection=map_h2_has, receiver=out_sig_has)
        sem_net_autodiff.add_projection(sender=h2, projection=map_h2_can, receiver=out_sig_can)
        # INPUTS & OUTPUTS FOR SEMANTIC NET:

        nouns = ['oak', 'pine', 'rose', 'daisy', 'canary', 'robin', 'salmon', 'sunfish']
        relations = ['is', 'has', 'can']
        is_list = ['living', 'living thing', 'plant', 'animal', 'tree', 'flower', 'bird', 'fish', 'big', 'green', 'red',
                   'yellow']
        has_list = ['roots', 'leaves', 'bark', 'branches', 'skin', 'feathers', 'wings', 'gills', 'scales']
        can_list = ['grow', 'move', 'swim', 'fly', 'breathe', 'breathe underwater', 'breathe air', 'walk', 'photosynthesize']

        nouns_input = np.identity(len(nouns))

        rels_input = np.identity(len(relations))

        truth_nouns = np.identity(len(nouns))

        truth_is = np.zeros((len(nouns), len(is_list)))

        truth_is[0, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0]
        truth_is[1, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0]
        truth_is[2, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
        truth_is[3, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
        truth_is[4, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[5, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[6, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0]
        truth_is[7, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0]

        truth_has = np.zeros((len(nouns), len(has_list)))

        truth_has[0, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[1, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[2, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[3, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[4, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[5, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[6, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]
        truth_has[7, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]

        truth_can = np.zeros((len(nouns), len(can_list)))

        truth_can[0, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[1, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[2, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[3, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[4, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[5, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[6, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]
        truth_can[7, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]

        # SETTING UP DICTIONARY OF INPUTS/OUTPUTS FOR SEMANTIC NET

        inputs_dict = {}
        inputs_dict[nouns_in] = []
        inputs_dict[rels_in] = []

        targets_dict = {}
        targets_dict[out_sig_I] = []
        targets_dict[out_sig_is] = []
        targets_dict[out_sig_has] = []
        targets_dict[out_sig_can] = []

        for i in range(len(nouns)):
            for j in range(len(relations)):
                inputs_dict[nouns_in].append(nouns_input[i])
                inputs_dict[rels_in].append(rels_input[j])
                targets_dict[out_sig_I].append(truth_nouns[i])
                targets_dict[out_sig_is].append(truth_is[i])
                targets_dict[out_sig_has].append(truth_has[i])
                targets_dict[out_sig_can].append(truth_can[i])

        inputs_dict_comp = {}
        inputs_dict_comp[nouns_in_comp] = inputs_dict[nouns_in]
        inputs_dict_comp[rels_in_comp] = inputs_dict[rels_in]

        sem_net_autodiff.run(inputs=inputs_dict)

        # TRAIN AUTODIFFCOMPOSITION
        def g_f():
            yield {"inputs": inputs_dict,
                   "targets": targets_dict,
                   "epochs": eps}
        g = g_f()
        sem_net_autodiff.learn(inputs=g_f, execution_mode=pnl.ExecutionMode.PyTorch)

        # SET UP COMPOSITION
        sem_net_comp = Composition()

        backprop_pathway = sem_net_comp.add_backpropagation_learning_pathway(
            pathway=[
                nouns_in_comp,
                map_nouns_h1_comp,
                h1_comp,
                map_h1_h2_comp,
                h2_comp,
                map_h2_I_comp,
                out_sig_I_comp
            ],
            learning_rate=0.5
        )
        inputs_dict_comp[backprop_pathway.target] = targets_dict[out_sig_I]

        backprop_pathway = sem_net_comp.add_backpropagation_learning_pathway(
            pathway=[
                rels_in_comp,
                map_rels_h2_comp,
                h2_comp,
                map_h2_is_comp,
                out_sig_is_comp
            ],
            learning_rate=0.5
        )
        inputs_dict_comp[backprop_pathway.target] = targets_dict[out_sig_is]

        backprop_pathway = sem_net_comp.add_backpropagation_learning_pathway(
            pathway=[
                h2_comp,
                map_h2_has_comp,
                out_sig_has_comp
            ],
            learning_rate=0.5
        )
        inputs_dict_comp[backprop_pathway.target] = targets_dict[out_sig_has]

        backprop_pathway = sem_net_comp.add_backpropagation_learning_pathway(
            pathway=[
                h2_comp,
                map_h2_can_comp,
                out_sig_can_comp
            ],
            learning_rate=0.5
        )
        inputs_dict_comp[backprop_pathway.target] = targets_dict[out_sig_can]

        # TRAIN COMPOSITION
        sem_net_comp.learn(inputs=inputs_dict_comp,
                           num_trials=(len(inputs_dict_comp[nouns_in_comp]) * eps))

        # CHECK THAT PARAMETERS FOR AUTODIFFCOMPOSITION, COMPOSITION ARE SAME

        np.testing.assert_allclose(map_nouns_h1.parameters.matrix.get(sem_net_autodiff),
                           map_nouns_h1_comp.get_mod_matrix(sem_net_comp))
        np.testing.assert_allclose(map_rels_h2.parameters.matrix.get(sem_net_autodiff),
                           map_rels_h2_comp.get_mod_matrix(sem_net_comp))
        np.testing.assert_allclose(map_h1_h2.parameters.matrix.get(sem_net_autodiff),
                           map_h1_h2_comp.get_mod_matrix(sem_net_comp))
        np.testing.assert_allclose(map_h2_I.parameters.matrix.get(sem_net_autodiff),
                           map_h2_I_comp.get_mod_matrix(sem_net_comp))
        np.testing.assert_allclose(map_h2_is.parameters.matrix.get(sem_net_autodiff),
                           map_h2_is_comp.get_mod_matrix(sem_net_comp))
        np.testing.assert_allclose(map_h2_has.parameters.matrix.get(sem_net_autodiff),
                           map_h2_has_comp.get_mod_matrix(sem_net_comp))
        np.testing.assert_allclose(map_h2_can.parameters.matrix.get(sem_net_autodiff),
                           map_h2_can_comp.get_mod_matrix(sem_net_comp))


@pytest.mark.pytorch
@pytest.mark.acnested
@pytest.mark.composition
class TestNestedNoLearning:

    @pytest.mark.parametrize(
        'num_epochs, learning_rate, patience, min_delta', [
            (400, 4, 10, .00001),
        ]
    )
    def test_xor_nested_train_then_no_train(self, num_epochs, learning_rate,
                                            patience, min_delta, autodiff_mode):
        # the inputs we will provide to the model
        xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])

        # the outputs we wish to see from the model
        xor_targets = np.array([[0], [1], [1], [0]])

        # -----------------------------------------------------------------

        xor_in = pnl.TransferMechanism(name='xor_in',
                                       default_variable=np.zeros(2))

        xor_hid = pnl.TransferMechanism(name='xor_hid',
                                        default_variable=np.zeros(10),
                                        function=Logistic())

        xor_out = pnl.TransferMechanism(name='xor_out',
                                        default_variable=np.zeros(1),
                                        function=Logistic())

        hid_map = pnl.MappingProjection(name='input_to_hidden',
                                        matrix=np.random.randn(2, 10) * 0.1,
                                        sender=xor_in,
                                        receiver=xor_hid)

        out_map = pnl.MappingProjection(name='hidden_to_output',
                                        matrix=np.random.randn(10, 1) * 0.1,
                                        sender=xor_hid,
                                        receiver=xor_out)

        # -----------------------------------------------------------------

        xor_autodiff = AutodiffComposition(
            learning_rate=learning_rate,
        )

        xor_autodiff.add_node(xor_in)
        xor_autodiff.add_node(xor_hid)
        xor_autodiff.add_node(xor_out)

        xor_autodiff.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor_autodiff.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # -----------------------------------------------------------------

        no_training_input_dict = {xor_in: xor_inputs}
        input_dict = {'inputs': {xor_in: xor_inputs }, 'targets': {xor_out: xor_targets}, 'epochs': num_epochs}

        parentComposition = pnl.Composition()
        parentComposition.add_node(xor_autodiff)

        input = {xor_autodiff: input_dict}
        no_training_input = {xor_autodiff: no_training_input_dict}

        learning_context = Context()
        xor_autodiff.learn(inputs=input_dict, execution_mode=autodiff_mode, epochs=num_epochs, context=learning_context, patience=patience, min_delta=min_delta)
        result1 = np.array(xor_autodiff.learning_results).flatten()
        np.testing.assert_allclose(result1, np.array(xor_targets).flatten(), atol=0.1)
        result2 = parentComposition.run(inputs=no_training_input, execution_mode=autodiff_mode, context=learning_context)

        np.testing.assert_allclose(result2, [[0]], atol=0.1)

    @pytest.mark.parametrize(
        'num_epochs, learning_rate, patience, min_delta', [
            (400, 4, 10, .00001),
        ]
    )
    @pytest.mark.llvm_not_implemented
    def test_xor_nested_no_train_then_train(self, num_epochs, learning_rate, patience, min_delta, autodiff_mode):
        xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
        xor_targets = np.array([[0], [1], [1], [0]])

        # -----------------------------------------------------------------

        xor_in = pnl.TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = pnl.TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = pnl.TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        hid_map = pnl.MappingProjection(name='input_to_hidden',
                                        matrix=np.random.randn(2, 10) * 0.1,
                                        sender=xor_in,
                                        receiver=xor_hid)

        out_map = pnl.MappingProjection(name='hidden_to_output',
                                        matrix=np.random.randn(10, 1) * 0.1,
                                        sender=xor_hid,
                                        receiver=xor_out)

        # -----------------------------------------------------------------

        xor_autodiff = AutodiffComposition(learning_rate=learning_rate)

        xor_autodiff.add_node(xor_in)
        xor_autodiff.add_node(xor_hid)
        xor_autodiff.add_node(xor_out)

        xor_autodiff.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor_autodiff.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # -----------------------------------------------------------------

        no_training_input_dict = {xor_in: xor_inputs}
        input_dict = {'inputs': {xor_in: xor_inputs}, 'targets': {xor_out: xor_targets}, 'epochs': num_epochs}

        parentComposition = pnl.Composition()
        parentComposition.add_node(xor_autodiff)
        input = {xor_autodiff: input_dict}
        no_training_input = {xor_autodiff: no_training_input_dict}

        learning_context = Context()

        result1 = xor_autodiff.run(inputs=input[xor_autodiff]['inputs'], execution_mode=autodiff_mode, context=learning_context)
        xor_autodiff.learn(inputs=input_dict, execution_mode=autodiff_mode, context=learning_context, patience=patience, min_delta=min_delta)
        result2 = parentComposition.run(inputs=no_training_input, execution_mode=autodiff_mode, context=learning_context)

        np.testing.assert_allclose(result2, [[0]], atol=0.1)

    # CW 12/21/18: Test is failing due to bugs, will fix later
    # @pytest.mark.parametrize(
    #     'num_epochs, learning_rate, patience, min_delta', [
    #         (2000, 4, 10, .00001),
    #     ]
    # )
    # def test_xor_nest_not_origin_after_train(self, num_epochs, learning_rate, patience, min_delta):
    #     xor_inputs = np.array(  # the inputs we will provide to the model
    #         [[0, 0],
    #          [0, 1],
    #          [1, 0],
    #          [1, 1]])
    #
    #     xor_targets = np.array(  # the outputs we wish to see from the model
    #         [[0],
    #          [1],
    #          [1],
    #          [0]])
    #
    #     # -----------------------------------------------------------------
    #
    #     xor_in = pnl.TransferMechanism(name='xor_in',
    #                                    default_variable=np.zeros(2))
    #
    #     xor_hid = pnl.TransferMechanism(name='xor_hid',
    #                                     default_variable=np.zeros(10),
    #                                     function=pnl.core.components.functions.transferfunctions.Logistic())
    #
    #     xor_out = pnl.TransferMechanism(name='xor_out',
    #                                     default_variable=np.zeros(1),
    #                                     function=pnl.core.components.functions.transferfunctions.Logistic())
    #
    #     hid_map = pnl.MappingProjection(name='input_to_hidden',
    #                                     matrix=np.random.randn(2, 10) * 0.1,
    #                                     sender=xor_in,
    #                                     receiver=xor_hid)
    #
    #     out_map = pnl.MappingProjection(name='hidden_to_output',
    #                                     matrix=np.random.randn(10, 1) * 0.1,
    #                                     sender=xor_hid,
    #                                     receiver=xor_out)
    #
    #     # -----------------------------------------------------------------
    #
    #     xor_autodiff = AutodiffComposition(
    #         patience=patience,
    #         min_delta=min_delta,
    #         learning_rate=learning_rate,
    #         randomize=False,
    #         learning_enabled=True
    #     )
    #
    #     xor_autodiff.add_node(xor_in)
    #     xor_autodiff.add_node(xor_hid)
    #     xor_autodiff.add_node(xor_out)
    #
    #     xor_autodiff.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
    #     xor_autodiff.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)
    #
    #     # -----------------------------------------------------------------
    #
    #     input_dict = {'inputs': {xor_in: xor_inputs}, 'targets': {xor_out: xor_targets}, 'epochs': num_epochs}
    #     xor_autodiff.run(inputs = input_dict)
    #     myTransfer = pnl.TransferMechanism(input_shapes = 2)
    #     myMappingProj = pnl.MappingProjection(sender = myTransfer, receiver = xor_autodiff)
    #
    #     no_training_input_dict = {xor_in: xor_inputs}
    #
    #     parentComposition = pnl.Composition()
    #     parentComposition.add_node(myTransfer)
    #     parentComposition.add_node(xor_autodiff)
    #     parentComposition.add_projection(myMappingProj, sender=myTransfer, receiver=xor_autodiff)
    #     xor_autodiff.learning_enabled = False
    #
    #     no_training_input = {myTransfer: no_training_input_dict}
    #
    #     result = parentComposition.run(inputs=no_training_input)
    #
    #     np.testing.assert_allclose(result, [[0]], atol=0.1)

    @pytest.mark.parametrize(
        'eps, opt', [
            (1, 'sgd'),
        ]
    )
    def test_semantic_net_nested(self, eps, opt, autodiff_mode):

        # SET UP MECHANISMS FOR SEMANTIC NET:

        nouns_in = TransferMechanism(name="nouns_input",
                                     default_variable=np.zeros(8))

        rels_in = TransferMechanism(name="rels_input",
                                    default_variable=np.zeros(3))

        h1 = TransferMechanism(name="hidden_nouns",
                               default_variable=np.zeros(8),
                               function=Logistic())

        h2 = TransferMechanism(name="hidden_mixed",
                               default_variable=np.zeros(15),
                               function=Logistic())

        out_sig_I = TransferMechanism(name="sig_outs_I",
                                      default_variable=np.zeros(8),
                                      function=Logistic())

        out_sig_is = TransferMechanism(name="sig_outs_is",
                                       default_variable=np.zeros(12),
                                       function=Logistic())

        out_sig_has = TransferMechanism(name="sig_outs_has",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        out_sig_can = TransferMechanism(name="sig_outs_can",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        # SET UP MECHANISMS FOR SYSTEM

        nouns_in_sys = TransferMechanism(name="nouns_input_sys",
                                         default_variable=np.zeros(8))

        rels_in_sys = TransferMechanism(name="rels_input_sys",
                                        default_variable=np.zeros(3))

        h1_sys = TransferMechanism(name="hidden_nouns_sys",
                                   default_variable=np.zeros(8),
                                   function=Logistic())

        h2_sys = TransferMechanism(name="hidden_mixed_sys",
                                   default_variable=np.zeros(15),
                                   function=Logistic())

        out_sig_I_sys = TransferMechanism(name="sig_outs_I_sys",
                                          default_variable=np.zeros(8),
                                          function=Logistic())

        out_sig_is_sys = TransferMechanism(name="sig_outs_is_sys",
                                           default_variable=np.zeros(12),
                                           function=Logistic())

        out_sig_has_sys = TransferMechanism(name="sig_outs_has_sys",
                                            default_variable=np.zeros(9),
                                            function=Logistic())

        out_sig_can_sys = TransferMechanism(name="sig_outs_can_sys",
                                            default_variable=np.zeros(9),
                                            function=Logistic())

        # SET UP PROJECTIONS FOR SEMANTIC NET

        map_nouns_h1 = MappingProjection(matrix=np.random.rand(8,8),
                                 name="map_nouns_h1",
                                 sender=nouns_in,
                                 receiver=h1)

        map_rels_h2 = MappingProjection(matrix=np.random.rand(3,15),
                                    name="map_relh2",
                                    sender=rels_in,
                                    receiver=h2)

        map_h1_h2 = MappingProjection(matrix=np.random.rand(8,15),
                                    name="map_h1_h2",
                                    sender=h1,
                                    receiver=h2)

        map_h2_I = MappingProjection(matrix=np.random.rand(15,8),
                                    name="map_h2_I",
                                    sender=h2,
                                    receiver=out_sig_I)

        map_h2_is = MappingProjection(matrix=np.random.rand(15,12),
                                    name="map_h2_is",
                                    sender=h2,
                                    receiver=out_sig_is)

        map_h2_has = MappingProjection(matrix=np.random.rand(15,9),
                                    name="map_h2_has",
                                    sender=h2,
                                    receiver=out_sig_has)

        map_h2_can = MappingProjection(matrix=np.random.rand(15,9),
                                    name="map_h2_can",
                                    sender=h2,
                                    receiver=out_sig_can)

        # SET UP PROJECTIONS FOR SYSTEM

        map_nouns_h1_sys = MappingProjection(matrix=map_nouns_h1.matrix.base.copy(),
                                             name="map_nouns_h1_sys",
                                             sender=nouns_in_sys,
                                             receiver=h1_sys)

        map_rels_h2_sys = MappingProjection(matrix=map_rels_h2.matrix.base.copy(),
                                        name="map_relh2_sys",
                                        sender=rels_in_sys,
                                        receiver=h2_sys)

        map_h1_h2_sys = MappingProjection(matrix=map_h1_h2.matrix.base.copy(),
                                          name="map_h1_h2_sys",
                                          sender=h1_sys,
                                          receiver=h2_sys)

        map_h2_I_sys = MappingProjection(matrix=map_h2_I.matrix.base.copy(),
                                         name="map_h2_I_sys",
                                         sender=h2_sys,
                                         receiver=out_sig_I_sys)

        map_h2_is_sys = MappingProjection(matrix=map_h2_is.matrix.base.copy(),
                                          name="map_h2_is_sys",
                                          sender=h2_sys,
                                          receiver=out_sig_is_sys)

        map_h2_has_sys = MappingProjection(matrix=map_h2_has.matrix.base.copy(),
                                           name="map_h2_has_sys",
                                           sender=h2_sys,
                                           receiver=out_sig_has_sys)

        map_h2_can_sys = MappingProjection(matrix=map_h2_can.matrix.base.copy(),
                                           name="map_h2_can_sys",
                                           sender=h2_sys,
                                           receiver=out_sig_can_sys)

        # SET UP COMPOSITION FOR SEMANTIC NET

        sem_net = AutodiffComposition(learning_rate=0.5,
                                      optimizer_type=opt)

        sem_net.add_node(nouns_in)
        sem_net.add_node(rels_in)
        sem_net.add_node(h1)
        sem_net.add_node(h2)
        sem_net.add_node(out_sig_I)
        sem_net.add_node(out_sig_is)
        sem_net.add_node(out_sig_has)
        sem_net.add_node(out_sig_can)

        sem_net.add_projection(sender=nouns_in, projection=map_nouns_h1, receiver=h1)
        sem_net.add_projection(sender=rels_in, projection=map_rels_h2, receiver=h2)
        sem_net.add_projection(sender=h1, projection=map_h1_h2, receiver=h2)
        sem_net.add_projection(sender=h2, projection=map_h2_I, receiver=out_sig_I)
        sem_net.add_projection(sender=h2, projection=map_h2_is, receiver=out_sig_is)
        sem_net.add_projection(sender=h2, projection=map_h2_has, receiver=out_sig_has)
        sem_net.add_projection(sender=h2, projection=map_h2_can, receiver=out_sig_can)

        # INPUTS & OUTPUTS FOR SEMANTIC NET:

        nouns = ['oak', 'pine', 'rose', 'daisy', 'canary', 'robin', 'salmon', 'sunfish']
        relations = ['is', 'has', 'can']
        is_list = ['living', 'living thing', 'plant', 'animal', 'tree', 'flower', 'bird', 'fish', 'big', 'green', 'red',
                   'yellow']
        has_list = ['roots', 'leaves', 'bark', 'branches', 'skin', 'feathers', 'wings', 'gills', 'scales']
        can_list = ['grow', 'move', 'swim', 'fly', 'breathe', 'breathe underwater', 'breathe air', 'walk', 'photosynthesize']

        nouns_input = np.identity(len(nouns))

        rels_input = np.identity(len(relations))

        truth_nouns = np.identity(len(nouns))

        truth_is = np.zeros((len(nouns), len(is_list)))

        truth_is[0, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0]
        truth_is[1, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0]
        truth_is[2, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
        truth_is[3, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
        truth_is[4, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[5, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[6, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0]
        truth_is[7, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0]

        truth_has = np.zeros((len(nouns), len(has_list)))

        truth_has[0, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[1, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[2, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[3, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[4, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[5, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[6, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]
        truth_has[7, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]

        truth_can = np.zeros((len(nouns), len(can_list)))

        truth_can[0, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[1, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[2, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[3, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[4, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[5, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[6, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]
        truth_can[7, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]

        # SETTING UP DICTIONARY OF INPUTS/OUTPUTS FOR SEMANTIC NET

        inputs_dict = {}
        inputs_dict[nouns_in] = []
        inputs_dict[rels_in] = []

        targets_dict = {}
        targets_dict[out_sig_I] = []
        targets_dict[out_sig_is] = []
        targets_dict[out_sig_has] = []
        targets_dict[out_sig_can] = []

        for i in range(len(nouns)):
            for j in range(len(relations)):
                inputs_dict[nouns_in].append(nouns_input[i])
                inputs_dict[rels_in].append(rels_input[j])
                targets_dict[out_sig_I].append(truth_nouns[i])
                targets_dict[out_sig_is].append(truth_is[i])
                targets_dict[out_sig_has].append(truth_has[i])
                targets_dict[out_sig_can].append(truth_can[i])

        # TRAIN COMPOSITION
        input_dict = {"inputs": inputs_dict,
                      "targets": targets_dict,
                      "epochs": eps}

        parentComposition = pnl.Composition()
        parentComposition.add_node(sem_net)

        input = {sem_net: input_dict}
        no_training_input = {sem_net: inputs_dict.copy()}

        sem_net.learn(inputs=input_dict, execution_mode=autodiff_mode)

        parentComposition.run(inputs=no_training_input)


@pytest.mark.pytorch
@pytest.mark.acnested
@pytest.mark.acidenticalness
@pytest.mark.composition
class TestNestedLearning:
    """Test learning with nested AutodiffCompositions benchmarked against unnested versions using Composition"""

    @pytest.fixture
    def nodes_for_testing_nested_comps(self):
        input_nodes = [pnl.ProcessingMechanism(name='input_1', input_shapes=2),
                       pnl.ProcessingMechanism(name='input_2', input_shapes=3),
                       pnl.ProcessingMechanism(name='input_3', input_shapes=3)]
        hidden_nodes = [pnl.ProcessingMechanism(name='hidden_1', input_shapes=3),
                        pnl.ProcessingMechanism(name='hidden_2', input_shapes=4),
                        pnl.ProcessingMechanism(name='hidden_3', input_shapes=5),
                        pnl.ProcessingMechanism(name='hidden_4', input_shapes=6)]
        output_nodes = [pnl.ProcessingMechanism(name='output_1', input_shapes=3),
                        pnl.ProcessingMechanism(name='output_2', input_shapes=5)]
        def _get_nodes(num_input_nodes, num_hidden_nodes, num_output_nodes):
            return (input_nodes[0:num_input_nodes],
                    hidden_nodes[0:num_hidden_nodes],
                    output_nodes[0:num_output_nodes])
        return _get_nodes

    @pytest.fixture
    def execute_learning(self):
        def _execute_learning(comp_type, execution_mode, pathways, inputs, learning_rate=.01, num_trials=5):

            if comp_type == 'composition':
                comp = Composition(name='comp')
                for p in pathways:
                    comp.add_backpropagation_learning_pathway(p)

            elif comp_type == 'autodiff':
                comp = AutodiffComposition(pathways, name='autodiff_comp')

            else:
                assert False, f'Invalid comp_type: {comp_type}'

            comp.learn(inputs=inputs,
                       learning_rate = learning_rate,
                       num_trials=num_trials,
                       execution_mode=execution_mode)
            return comp.results
        return _execute_learning

    def test_warning_for_no_learning_in_solo_nested_comp(self):
        with pytest.warns(UserWarning) as warning:
            inner = pnl.AutodiffComposition([pnl.ProcessingMechanism()])
            outer = pnl.AutodiffComposition([inner])
            outer._build_pytorch_representation()
        assert (f"'autodiff_composition-1' contains no Projections, so it has no params for Pytorch to learn."
                in repr(warning[0].message.args[0]))

    def test_1_nested_hidden(self, nodes_for_testing_nested_comps, execute_learning):
        nodes = nodes_for_testing_nested_comps(1, 1, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested_hidden_1 = AutodiffComposition(name='nested', nodes=[hidden_nodes[0]])
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[input_nodes[0], nested_hidden_1, output_nodes[0]],
                                            inputs=inputs)

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[[input_nodes[0],hidden_nodes[0], output_nodes[0]]],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results,  rtol=1e-8, atol=1e-8)

    def test_2_sequential_nested_hidden(self, nodes_for_testing_nested_comps, execute_learning):
        nodes = nodes_for_testing_nested_comps(1, 2, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested_hiddens = AutodiffComposition(name='nested', pathways=[hidden_nodes[0],
                                                                      hidden_nodes[1]])
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[input_nodes[0], nested_hiddens, output_nodes[0]],
                                            inputs=inputs)

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[[input_nodes[0],
                                                   hidden_nodes[0],
                                                   hidden_nodes[1],
                                                   output_nodes[0]]],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_2_level_nested(self, nodes_for_testing_nested_comps, execute_learning):

        nodes = nodes_for_testing_nested_comps(1, 4, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested_b = AutodiffComposition(name='nested_b', pathways=[hidden_nodes[1],
                                                                  hidden_nodes[2]])
        nested_a = AutodiffComposition(name='nested_a', pathways=[hidden_nodes[0],
                                                                  nested_b,
                                                                  hidden_nodes[3]])
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[input_nodes[0], nested_a, output_nodes[0]],
                                            inputs=inputs)

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[[input_nodes[0],
                                                   hidden_nodes[0],
                                                   hidden_nodes[1],
                                                   hidden_nodes[2],
                                                   hidden_nodes[3],
                                                   output_nodes[0]]],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_1_input_to_1_nested_hidden_with_2_output_ports(self, nodes_for_testing_nested_comps, execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(1, 1, 2)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}
        hidden_with_two_output_ports = pnl.ProcessingMechanism(input_shapes=3, output_ports=['FIRST', 'SECOND'])

        nested = AutodiffComposition([hidden_nodes[0], hidden_with_two_output_ports], name='nested')
        pathway_a = [input_nodes[0],
                     nested,
                     MappingProjection(hidden_with_two_output_ports.output_ports[0], output_nodes[0]),
                     output_nodes[0]]
        pathway_b = [input_nodes[0],
                     nested,
                     MappingProjection(hidden_with_two_output_ports.output_ports[1], output_nodes[1]),
                     output_nodes[1]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actually implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0], hidden_nodes[0], hidden_with_two_output_ports, output_nodes[0]]
        pathway_b = [input_nodes[0], hidden_nodes[0], hidden_with_two_output_ports, output_nodes[1]]

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    def test_1_input_to_1_nested_hidden_one_to_many_2_outputs(self, nodes_for_testing_nested_comps, execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(1, 1, 2)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition([hidden_nodes[0]],name='nested')
        pathway_a = [input_nodes[0], nested, output_nodes[0]]
        pathway_b = [input_nodes[0], nested, output_nodes[1]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actualy implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0], hidden_nodes[0], output_nodes[0]]
        pathway_b = [input_nodes[0], hidden_nodes[0], output_nodes[1]]

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    def test_1_input_to_2_nested_hidden_to_2_outputs(self, nodes_for_testing_nested_comps, execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(1, 2, 2)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition(nodes=[hidden_nodes[0], hidden_nodes[1]], name='nested')
        pathway_a = [input_nodes[0],
                     MappingProjection(input_nodes[0],hidden_nodes[0]),
                     nested,
                     MappingProjection(hidden_nodes[0],output_nodes[0]),
                     output_nodes[0]]
        pathway_b = [input_nodes[0],
                     MappingProjection(input_nodes[0], hidden_nodes[1]),
                     nested,
                     MappingProjection(hidden_nodes[1],output_nodes[1]),
                     output_nodes[1]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actualy implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0], hidden_nodes[0], output_nodes[0]]
        pathway_b = [input_nodes[0], hidden_nodes[1], output_nodes[1]]

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    def test_2_inputs_to_2_input_ports_of_single_nested_hidden(self, nodes_for_testing_nested_comps, execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(2, 0, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        hidden_with_2_inputs = pnl.ProcessingMechanism(name='hidden_x', input_shapes=(3, 3), function=pnl.LinearCombination)

        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition([hidden_with_2_inputs],name='nested')
        pathway_a = [input_nodes[0],
                     MappingProjection(input_nodes[0], hidden_with_2_inputs.input_ports[0]),
                     nested,
                     output_nodes[0]]
        pathway_b = [input_nodes[1],
                     MappingProjection(input_nodes[1], hidden_with_2_inputs.input_ports[1]),
                     nested,
                     output_nodes[0]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actualy implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0],
                     # MappingProjection(input_nodes[0], hidden_with_2_inputs.input_ports[0]),
                     hidden_with_2_inputs,
                     output_nodes[0]]
        pathway_b = [input_nodes[1],
                     # MappingProjection(input_nodes[1], hidden_with_2_inputs.input_ports[1]),
                     hidden_with_2_inputs,
                     output_nodes[0]]

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_2_inputs_one_to_one_to_2_nested_hidden_to_1_hidden(self, nodes_for_testing_nested_comps,execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(2, 3, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition([[hidden_nodes[0], hidden_nodes[2]],
                                      [hidden_nodes[1], hidden_nodes[2]]],
                                     name='nested')
        pathway_a = [input_nodes[0], MappingProjection(input_nodes[0], hidden_nodes[0]), nested, output_nodes[0]]
        pathway_b = [input_nodes[1], MappingProjection(input_nodes[1], hidden_nodes[1]), nested, output_nodes[0]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actualy implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0], hidden_nodes[0], hidden_nodes[2], output_nodes[0]]
        pathway_b = [input_nodes[1], hidden_nodes[1], hidden_nodes[2], output_nodes[0]]
        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_2_inputs_all_to_all_to_2_nested_hidden_to_1_hidden(self, nodes_for_testing_nested_comps, execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(2, 3, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition([[hidden_nodes[0], hidden_nodes[2]],
                                      [hidden_nodes[1], hidden_nodes[2]]],
                                      name='nested')
        pathway_a = [input_nodes[0], nested, output_nodes[0]]
        pathway_b = [input_nodes[1], nested, output_nodes[0]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actualy implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0], hidden_nodes[0], hidden_nodes[2], output_nodes[0]]
        pathway_b = [input_nodes[0], hidden_nodes[1], hidden_nodes[2], output_nodes[0]]
        pathway_c = [input_nodes[1], hidden_nodes[0], hidden_nodes[2], output_nodes[0]]
        pathway_d = [input_nodes[1], hidden_nodes[1], hidden_nodes[2], output_nodes[0]]
        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b, pathway_c, pathway_d],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_2_inputs_to_2_parallel_nested_and_output_pathways(self, nodes_for_testing_nested_comps, execute_learning):
        # Note: inputs provided for only one of the INPUT Nodes; other uses default inputs

        nodes = nodes_for_testing_nested_comps(2, 2, 2)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition(nodes=[hidden_nodes[0], hidden_nodes[1]], name='nested')
        pathway_a = [input_nodes[0],
                     MappingProjection(input_nodes[0], hidden_nodes[0]),
                     nested,
                     MappingProjection(hidden_nodes[0], output_nodes[0]),
                     output_nodes[0]]
        pathway_b = [input_nodes[1],
                     MappingProjection(input_nodes[1], hidden_nodes[1]),
                     nested,
                     MappingProjection(hidden_nodes[1], output_nodes[1]),
                     output_nodes[1]]
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        # Note:
        #  the MappingProjections are not needed, as they were actualy implemented in the autodiff version
        #  in infer_backpropagation_learning_pathways() (when flattening the nessted Composition)
        pathway_a = [input_nodes[0], hidden_nodes[0], output_nodes[0]]
        pathway_b = [input_nodes[1], hidden_nodes[1], output_nodes[1]]
        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    def test_direct_input_to_nested(self, nodes_for_testing_nested_comps, execute_learning):
        nodes = nodes_for_testing_nested_comps(1, 1, 1)
        input_nodes, hidden_nodes, output_nodes = nodes

        nested = AutodiffComposition([input_nodes[0], hidden_nodes[0]],name='nested')
        auto_diff_inputs = {nested:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}
        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[nested, output_nodes[0]],
                                            inputs=auto_diff_inputs)

        comp_inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}
        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[[input_nodes[0],hidden_nodes[0], output_nodes[0]]],
                                        inputs=comp_inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_asymmetric_inputs_to_nested_nodes_one_direct_one_via_node_in_outer_comp(
            self,
            nodes_for_testing_nested_comps,
            execute_learning):
        nodes = nodes_for_testing_nested_comps(1, 2, 2)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition(nodes=[hidden_nodes[0],hidden_nodes[1]],name='nested')
        indirect = [input_nodes[0],
                  MappingProjection(input_nodes[0], hidden_nodes[0]),
                  nested,
                  MappingProjection(hidden_nodes[0], output_nodes[0]),
                  output_nodes[0]]
        direct = [(nested, pnl.NodeRole.INPUT),
                  MappingProjection(hidden_nodes[1], output_nodes[1]),
                  output_nodes[1]]

        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[indirect, direct],
                                            inputs=inputs)

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[[input_nodes[0],hidden_nodes[0], output_nodes[0]],
                                                  [hidden_nodes[1], output_nodes[1]]],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    def test_inputs_to_multiple_input_ports_and_INPUT_nodes(self, nodes_for_testing_nested_comps, execute_learning):
        """Test asymmetric inputs to INPUT Nodes of a nested Composition
        One direct input to one input_port of an INPUT Node of nested Composition,
        One indirect (via Node in outer Composition) input to the other input_port of the nested INPUT Node
        One direct input to a different INPUT Node of the nested Composition
        """
        nodes = nodes_for_testing_nested_comps(2, 1, 2)
        input_nodes, hidden_nodes, output_nodes = nodes
        hidden_node_x = pnl.ComparatorMechanism(name='hidden_node_x')

        nested = AutodiffComposition(nodes=[hidden_nodes[0],hidden_node_x],name='nested')

        inputs = {input_nodes[0]: np.array([[0, 0], [0, 1], [1, 0], [1, 1]]),
                  hidden_node_x.input_ports[0]: [[1], [1], [0], [0]]}

        direct_1 = [(nested, pnl.NodeRole.INPUT),
                  MappingProjection(hidden_nodes[0], output_nodes[0]),
                  output_nodes[0]]
        direct_2 = [(nested, pnl.NodeRole.INPUT),
                  MappingProjection(hidden_node_x, output_nodes[1]),
                  output_nodes[1]]
        indirect = [input_nodes[0],
                      MappingProjection(input_nodes[0], hidden_node_x.input_ports[1]),
                      nested]

        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[indirect, direct_1, direct_2],
                                            inputs=inputs)

        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[[hidden_nodes[0], output_nodes[0]],
                                                  [input_nodes[0],
                                                   # MappingProjection(input_nodes[0], hidden_node_x.input_ports[1]),
                                                   (hidden_node_x, pnl.NodeRole.INPUT),
                                                   output_nodes[1]]],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    # def test_output_from_multiple_output_ports_and_OUTPUT_nodes(self, nodes_for_testing_nested_comps, execute_learning):
    #     # FIX:  1) No input assigned to hidden_2d.input_port[1]
    #     #       2) Output from hidden_2d.output_port[0] erroneously projects to output_nodes[1]
    #     #          in addition to correctly projecting to output_nodes[0]
    #     nodes = nodes_for_testing_nested_comps(1, 1, 2)
    #     input_nodes, hidden_nodes, output_nodes = nodes
    #     inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}
    #
    #     hidden_2d = pnl.ProcessingMechanism(name='hidden 2d', input_shapes=(2,2))
    #     nested = AutodiffComposition(nodes = [hidden_nodes[0], hidden_2d], name='nested')
    #     pathway_a = [input_nodes[0],
    #                  MappingProjection(input_nodes[0], hidden_2d),
    #                  nested,
    #                  MappingProjection(hidden_2d.output_ports[0], output_nodes[0]),
    #                  output_nodes[0]]
    #     pathway_b = [input_nodes[0], MappingProjection(input_nodes[0], hidden_nodes[0]), nested, output_nodes[1]]
    #     # pathway_c = [nested,
    #     #              MappingProjection(hidden_2d.output_ports[1], output_nodes[1]),
    #     #              output_nodes[1]]
    #
    #     autodiff_results = execute_learning(comp_type='autodiff',
    #                                         execution_mode=pnl.ExecutionMode.PyTorch,
    #                                         pathways=[pathway_a, pathway_b],
    #                                         inputs=inputs)
    #
    #     pathway_a = [input_nodes[0], hidden_nodes[0]]
    #     pathway_b = [input_nodes[0], hidden_nodes[1]]
    #     comp_results = execute_learning(comp_type='composition',
    #                                     execution_mode=pnl.ExecutionMode.Python,
    #                                     pathways=[pathway_a, pathway_b],
    #                                     inputs=inputs)
    #
    #     for i in range(len(autodiff_results)):
    #         for j in range(len(autodiff_results[i])):
    #             np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])
    def test_2_direct_output_nodes_from_nested(self, nodes_for_testing_nested_comps, execute_learning):
        nodes = nodes_for_testing_nested_comps(1, 2, 0)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition(nodes = [hidden_nodes[0], hidden_nodes[1]], name='nested')
        pathway_a = [input_nodes[0], MappingProjection(input_nodes[0], hidden_nodes[0]), nested]
        pathway_b = [input_nodes[0], MappingProjection(input_nodes[0], hidden_nodes[1]), nested]

        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        pathway_a = [input_nodes[0], hidden_nodes[0]]
        pathway_b = [input_nodes[0], hidden_nodes[1]]
        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        for i in range(len(autodiff_results)):
            for j in range(len(autodiff_results[i])):
                np.testing.assert_allclose(comp_results[i][j], autodiff_results[i][j])

    def test_1_direct_and_1_ordinary_output_from_nested(self, nodes_for_testing_nested_comps, execute_learning):
        nodes = nodes_for_testing_nested_comps(1, 2, 1)
        input_nodes, hidden_nodes, output_nodes = nodes
        inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}

        nested = AutodiffComposition(nodes = [hidden_nodes[0], hidden_nodes[1]], name='nested')
        pathway_a = [input_nodes[0], MappingProjection(input_nodes[0], hidden_nodes[0]), nested]
        pathway_b = [input_nodes[0],
                     MappingProjection(input_nodes[0], hidden_nodes[1]),
                     nested,
                     MappingProjection(hidden_nodes[1]), # <- FIX FAILS
                     output_nodes[0]]

        autodiff_results = execute_learning(comp_type='autodiff',
                                            execution_mode=pnl.ExecutionMode.PyTorch,
                                            pathways=[pathway_a, pathway_b],
                                            inputs=inputs)

        pathway_a = [input_nodes[0], hidden_nodes[0]]
        pathway_b = [input_nodes[0], hidden_nodes[1], output_nodes[0]]
        comp_results = execute_learning(comp_type='composition',
                                        execution_mode=pnl.ExecutionMode.Python,
                                        pathways=[pathway_a, pathway_b],
                                        inputs=inputs)

        np.testing.assert_allclose(comp_results, autodiff_results)

    def test_nested_autodiff_learning_with_input_func(self):
        """Note: this uses the same Composition and results as test_learning/test_identicalness_of_input_types"""
        xor_in_func = TransferMechanism(name='xor_in',
                                        default_variable=np.zeros(2))

        xor_hid_func = TransferMechanism(name='xor_hid',
                                         default_variable=np.zeros(10),
                                         function=Logistic())

        xor_out_func = TransferMechanism(name='xor_out',
                                         default_variable=np.zeros(1),
                                         function=Logistic())

        nested = AutodiffComposition([xor_hid_func], learning_rate=.001, name='nested')
        xor_func = AutodiffComposition([xor_in_func,
                                        MappingProjection(sender=xor_in_func,
                                                          matrix=np.full((2, 10),.1),
                                                          receiver=xor_hid_func),
                                        nested,
                                        MappingProjection(sender=xor_hid_func,
                                                          matrix=np.full((10, 1),.1),
                                                          receiver=xor_out_func),
                                        # out_map_func,
                                        xor_out_func],
                                       learning_rate=.001,
                                       name='outer')

        xor_inputs_func = np.array([[0, 0],[0, 1],[1, 0],[1, 1]])
        xor_targets_func = np.array([[0],[1],[1],[0]])

        def get_inputs_auto_diff(idx):
            return {"inputs": {xor_in_func: xor_inputs_func[idx]},
                    "targets": {xor_out_func: xor_targets_func[idx]}}

        def get_inputs_comp(idx):
            return {xor_in_func: xor_inputs_func[idx]}

        def get_targets_comp(idx):
            return {xor_out_func: xor_targets_func[idx]}

        xor_func.learn(inputs=get_inputs_auto_diff, execution_mode=pnl.ExecutionMode.PyTorch)

        results = xor_func.results
        np.testing.assert_allclose(results, [[[0.62245933]],[[0.62813197]],[[0.6282438]],[[0.6341436]]])

    # def test_nested_with_diff_learning_rates(self, nodes_for_testing_nested_comps, execute_learning):
    #     nodes = nodes_for_testing_nested_comps(1, 0, 1)
    #     input_nodes, hidden_nodes, output_nodes = nodes
    #     inputs = {input_nodes[0]:np.array([[0, 0], [0, 1], [1, 0], [1, 1]])}
    #
    #     hidden_1 = pnl.ProcessingMechanism(name='hidden_1', input_shapes=3)
    #     nested_01 = AutodiffComposition(name='nested_01', nodes=[hidden_1], learning_rate=.01)
    #     autodiff_01_results = execute_learning(comp_type='autodiff',
    #                                           execution_mode=pnl.ExecutionMode.PyTorch,
    #                                           pathways=[input_nodes[0], nested_01, output_nodes[0]],
    #                                           inputs=inputs)
    #
    #     hidden_2 = pnl.ProcessingMechanism(name='hidden_2', input_shapes=3)
    #     nested_1 = AutodiffComposition(name='nested_2', nodes=[hidden_2], learning_rate=.1)
    #     autodiff_1_results = execute_learning(comp_type='autodiff',
    #                                           execution_mode=pnl.ExecutionMode.PyTorch,
    #                                           pathways=[input_nodes[0], nested_1, output_nodes[0]],
    #                                           inputs=inputs)
    #
    #     THIS SHOULD FAIL IF NESTED USED ITS OWN SPECIFIED LEARNING RATE:
    #     np.testing.assert_allclose(autodiff_01_results, autodiff_1_results)

    def test_error_for_running_nested_learning_in_Python_mode(self):
        input_mech = pnl.ProcessingMechanism(name='input_mech', input_shapes=2)
        hidden_mech = pnl.ProcessingMechanism(name='hidden_mech', input_shapes=2)
        output_mech = pnl.ProcessingMechanism(name='output_mech', input_shapes=2)

        # Test for error on learning if nested is Composition
        nested = pnl.Composition(name='nested', nodes=[hidden_mech])
        autodiff = pnl.AutodiffComposition([input_mech, nested, output_mech], name='comp')
        autodiff.run([0, 0])

        error_msg = (f"Unable execute learning for 'comp' because it contains nested Composition(s) "
                     f"that are not AutodiffCompositions: 'nested'.")
        with pytest.raises(AutodiffCompositionError) as error:
            autodiff.learn(inputs={input_mech: [[0, 0]]}, execution_mode=pnl.ExecutionMode.PyTorch)

        assert error_msg in str(error.value)

        # Test for error on learning if nested is AutodiffComposition but execution_mode is Python
        nested = pnl.AutodiffComposition(name='nested', nodes=[hidden_mech])
        autodiff = pnl.AutodiffComposition([input_mech, nested, output_mech], name='comp')
        autodiff.run([0, 0])

        error_msg = ("Unable to execute learning in Python mode for 'comp-1' "
                     "because it contains one or more nested Compositions: 'nested-1'.")
        with pytest.raises(AutodiffCompositionError) as error:
            autodiff.learn(inputs={input_mech: [[0, 0]]}, execution_mode=pnl.ExecutionMode.Python)

        assert error_msg in str(error.value)


# Names and objects for TestAutodiffMultipleOutput_ports
INPUT_A = 'input_A'
INPUT_B = 'input_B'
INPUT_C = 'input_C'
HIDDEN_A = 'hidden_A'
HIDDEN_B = 'hidden_B'
HIDDEN_C = 'hidden_C'
OUTPUT_A = 'output_A'
OUTPUT_B = 'output_B'
def nodes_for_testing_nested_comps(sizes):
    return {INPUT_A: pnl.ProcessingMechanism(name=INPUT_A, input_shapes=sizes.pop(INPUT_A, 2)),
            INPUT_B: pnl.ProcessingMechanism(name=INPUT_B, input_shapes=sizes.pop(INPUT_B, 2)),
            INPUT_C: pnl.ProcessingMechanism(name=INPUT_C, input_shapes=sizes.pop(INPUT_C, 2)),
            HIDDEN_A: pnl.ProcessingMechanism(name=HIDDEN_A, input_shapes=sizes.pop(HIDDEN_A, 2)),
            HIDDEN_B: pnl.ProcessingMechanism(name=HIDDEN_B, input_shapes=sizes.pop(HIDDEN_B, 2)),
            HIDDEN_C: pnl.ProcessingMechanism(name=HIDDEN_C, input_shapes=sizes.pop(HIDDEN_C, 2)),
            OUTPUT_A: pnl.ProcessingMechanism(name=OUTPUT_A, input_shapes=sizes.pop(OUTPUT_A, 2)),
            OUTPUT_B: pnl.ProcessingMechanism(name=OUTPUT_B, input_shapes=sizes.pop(OUTPUT_B, 2))}


@pytest.mark.pytorch
class TestAutodiffMultipleOutput_ports:

    def test_parallel_inputs_to_output_ports_converge_internal(self):

        # Autodiff: PARALLEL PATHWAYS USING SEPARATE INPUT AND OUTPUT PORTS ON hidden_1
        sizes = {HIDDEN_A: (2,3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input_A = nodes[INPUT_A]
        input_B = nodes[INPUT_B]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        output = nodes[OUTPUT_A]
        autodiff = pnl.AutodiffComposition(pathways=[[input_A,
                                                      pnl.MappingProjection(input_A, hidden_A.input_ports[0]),
                                                      hidden_A,
                                                      pnl.MappingProjection(hidden_A.output_ports[0],hidden_B),
                                                      hidden_B,
                                                      output],
                                                     [input_B,
                                                      pnl.MappingProjection(input_B, hidden_A.input_ports[1]),
                                                      hidden_A,
                                                      pnl.MappingProjection(hidden_A.output_ports[1],hidden_B),
                                                      hidden_B,
                                                      output]],
                                           name='autodiff')
        # autodiff.show_graph(show_all=True)
        result_autodiff_ports = autodiff.learn(inputs={input_A: [[0, 0], [0, 1], [1, 0], [1, 1]],
                                                       input_B: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # # Autodiff: PARALLEL PATHWAYS USING SEPARATE HIDDEN NODES (hidden_1 VS. hidden_2)
        sizes = {HIDDEN_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input_A = nodes[INPUT_A]
        input_B = nodes[INPUT_B]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        hidden_C = nodes[HIDDEN_C]
        output = nodes[OUTPUT_A]
        autodiff = pnl.AutodiffComposition(pathways=[[input_A, hidden_A, hidden_C, output],
                                                     [input_B, hidden_B, hidden_C, output]],
                                           name='autodiff')

        result_autodiff_nodes = autodiff.learn(inputs={input_A: [[0, 0], [0, 1], [1, 0], [1, 1]],
                                                       input_B: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # # Composition: PARALLEL PATHWAYS USING SEPARATE INPUT AND OUTPUT PORTS ON hidden_1
        sizes = {HIDDEN_A: (2,3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input_A = nodes[INPUT_A]
        input_B = nodes[INPUT_B]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        output = nodes[OUTPUT_A]
        comp = Composition(name='comp')
        comp.add_backpropagation_learning_pathway([input_A,
                                                   pnl.MappingProjection(input_A, hidden_A.input_ports[0]),
                                                   hidden_A,
                                                   pnl.MappingProjection(hidden_A.output_ports[0],hidden_B),
                                                   hidden_B,
                                                   output])
        comp.add_backpropagation_learning_pathway([input_B,
                                                   pnl.MappingProjection(input_B, hidden_A.input_ports[1]),
                                                   hidden_A,
                                                   pnl.MappingProjection(hidden_A.output_ports[1],hidden_B),
                                                   hidden_B,
                                                   output])
        result_comp_ports = comp.learn(inputs={input_A: [[0, 0], [0, 1], [1, 0], [1, 1]],
                                               input_B: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                       learning_rate=.01,
                                       epochs=3)

        # # Comp: PARALLEL PATHWAYS USING SEPARATE HIDDEN NODES (hidden_1 VS. hidden_2)
        sizes = {HIDDEN_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input_A = nodes[INPUT_A]
        input_B = nodes[INPUT_B]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        hidden_C = nodes[HIDDEN_C]
        output = nodes[OUTPUT_A]
        comp = Composition(name='comp')
        comp.add_backpropagation_learning_pathway([input_A,
                                                   hidden_A,
                                                   hidden_C,
                                                   output])
        comp.add_backpropagation_learning_pathway([input_B,
                                                   hidden_B,
                                                   hidden_C,
                                                   output])
        result_comp_nodes = comp.learn(inputs={input_A: [[0, 0], [0, 1], [1, 0], [1, 1]],
                                               input_B: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                       learning_rate=.01,
                                       epochs=3)

        expected = [[-0.14091702170015408, 0.11156579308015635]]
        np.testing.assert_allclose(result_autodiff_ports, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_autodiff_nodes, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_comp_ports, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_comp_nodes, expected, rtol=1e-8, atol=1e-8)

    def test_single_input_to_multiple_output_ports_converge_internal(self):

        # Autodiff: DIVERGENT PATHWAY USING SEPARATE INPUT AND OUTPUT PORTS ON hidden_1
        sizes = {HIDDEN_A: (2, 3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        output = nodes[OUTPUT_A]
        autodiff = pnl.AutodiffComposition(pathways=[[input,
                                                      pnl.MappingProjection(input, hidden_A.input_ports[0]),
                                                      hidden_A,
                                                      pnl.MappingProjection(hidden_A.output_ports[0],hidden_B),
                                                      hidden_B,
                                                      output],
                                                     [input,
                                                      pnl.MappingProjection(input, hidden_A.input_ports[1]),
                                                      hidden_A,
                                                      pnl.MappingProjection(hidden_A.output_ports[1],hidden_B),
                                                      hidden_B,
                                                      output]],
                                           name='autodiff')
        result_autodiff_ports = autodiff.learn(inputs={input: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # # Autodiff: DIVERGENT PATHWAY USING SEPARATE HIDDEN NODES (hidden_1 VS. hidden_2)
        sizes = {HIDDEN_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        hidden_C = nodes[HIDDEN_C]
        output = nodes[OUTPUT_A]
        autodiff = pnl.AutodiffComposition(pathways=[[input, hidden_A, hidden_C, output],
                                                     [input, hidden_B, hidden_C, output]],
                                           name='autodiff')
        result_autodiff_nodes = autodiff.learn(inputs={input: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # Composition: DIVERGENT PATHWAY USING SEPARATE INPUT AND OUTPUT PORTS ON hidden_1
        sizes = {HIDDEN_A: (2, 3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        output = nodes[OUTPUT_A]
        comp = Composition(name='comp')
        comp.add_backpropagation_learning_pathway([input,
                                                   pnl.MappingProjection(input, hidden_A.input_ports[0]),
                                                   hidden_A,
                                                   pnl.MappingProjection(hidden_A.output_ports[0],hidden_B),
                                                   hidden_B,
                                                   output])
        comp.add_backpropagation_learning_pathway([input,
                                                   pnl.MappingProjection(input, hidden_A.input_ports[1]),
                                                   hidden_A,
                                                   pnl.MappingProjection(hidden_A.output_ports[1],hidden_B),
                                                   hidden_B,
                                                   output])
        result_comp_ports = comp.learn(inputs={input: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                       learning_rate=.01,
                                       epochs=3)

        # Comp: DIVERGENT PATHWAY USING SEPARATE HIDDEN NODES (hidden_1 VS. hidden_2)
        sizes = {HIDDEN_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        hidden_C = nodes[HIDDEN_C]
        output = nodes[OUTPUT_A]
        comp = Composition(name='comp')
        comp.add_backpropagation_learning_pathway([input, hidden_A, hidden_C, output])
        comp.add_backpropagation_learning_pathway([input, hidden_B, hidden_C, output])
        result_comp_nodes = comp.learn(inputs={input: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                       learning_rate=.01,
                                       epochs=3)

        expected = [[-0.5448706019245989, -0.45743872720005285]]
        np.testing.assert_allclose(result_autodiff_ports, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_autodiff_nodes, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_comp_ports, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_comp_nodes, expected, rtol=1e-8, atol=1e-8)

    def test_single_input_to_multiple_output_ports_converge_on_OUTPUT_Node(self):

        # Autodiff: PARALLEL PATHWAYS USE SEPARATE INPUT AND OUTPUT PORTS ON hidden
        sizes = {HIDDEN_A: (2,3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden = nodes[HIDDEN_A]
        output = nodes[OUTPUT_A]
        autodiff = pnl.AutodiffComposition(pathways=[[input,
                                                  pnl.MappingProjection(input, hidden.input_ports[0]),
                                                  hidden,
                                                  pnl.MappingProjection(hidden.output_ports[0],output),
                                                  output],
                                                 [input,
                                                  pnl.MappingProjection(input, hidden.input_ports[1]),
                                                  hidden,
                                                  pnl.MappingProjection(hidden.output_ports[1],output),
                                                  output]                                         ],
                                       name='autodiff')
        result_autodiff_ports = autodiff.learn(inputs={input: [[0, 0], [0, 1], [1, 0], [1, 1]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # Autodiff: PARALLEL PATHWAYS USE SEPARATE HIDDEN NODES (hidden_1 VS. hidden_2)
        sizes = {HIDDEN_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        output = nodes[OUTPUT_A]
        autodiff = pnl.AutodiffComposition(pathways=[[input, hidden_A, output], [input, hidden_B, output]], name='autodiff')
        result_autodiff_nodes = autodiff.learn(inputs={input: [[0, 0], [0, 1], [1, 0], [1, 1]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # Comp: PARALLEL PATHWAYS USE SEPARATE HIDDEN NODES (hidden_1 VS. hidden_2)
        sizes = {HIDDEN_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden_A = nodes[HIDDEN_A]
        hidden_B = nodes[HIDDEN_B]
        output = nodes[OUTPUT_A]
        comp = Composition(name='comp')
        comp.add_backpropagation_learning_pathway([input, hidden_A, output])
        comp.add_backpropagation_learning_pathway([input, hidden_B, output])
        result_comp_nodes = comp.learn(inputs={input: [[0, 0], [0, 1], [1, 0], [1, 1]]},
                                       learning_rate=.01,
                                       epochs=3)

        # Comp: PARALLEL PATHWAYS USE SEPARATE INPUT AND OUTPUT PORTS ON hidden_c
        sizes = {HIDDEN_A: (2,3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input = nodes[INPUT_A]
        hidden = nodes[HIDDEN_A]
        output = nodes[OUTPUT_A]
        comp = Composition(name='comp')
        comp.add_backpropagation_learning_pathway([input,
                                                   pnl.MappingProjection(input, hidden.input_ports[0]),
                                                   hidden,
                                                   pnl.MappingProjection(hidden.output_ports[0],output),
                                                   output])
        comp.add_backpropagation_learning_pathway([input,
                                                   pnl.MappingProjection(input, hidden.input_ports[1]),
                                                   hidden,
                                                   pnl.MappingProjection(hidden.output_ports[1],output),
                                                   output])
        result_comp_ports = comp.learn(inputs={input: [[0, 0], [0, 1], [1, 0], [1, 1]]},
                                       learning_rate=.01,
                                       epochs=3)

        expected = [[3.3178720430554267, 3.3245710462077773]]
        np.testing.assert_allclose(result_autodiff_ports, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_autodiff_nodes, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_comp_ports, expected, rtol=1e-8, atol=1e-8)
        np.testing.assert_allclose(result_comp_nodes, expected, rtol=1e-8, atol=1e-8)

    def test_two_output_ports_on_OUTPUT_Node(self):

        # Autodiff: SEPARATE INPUT NODES TO ONE OUTPUT NODE WITH TWO OUTPUTPORTS
        sizes = {OUTPUT_A: (2,3)}
        nodes = nodes_for_testing_nested_comps(sizes)
        input_A = nodes[INPUT_A]
        input_B = nodes[INPUT_B]
        output = nodes[OUTPUT_A]
        autodiff = AutodiffComposition(pathways=[[input_A, MappingProjection(input_A, output.input_ports[0]), output],
                                                 [input_B, MappingProjection(input_B, output.input_ports[1]), output]],
                                       name='autodiff')
        result_autodiff_ports = autodiff.learn(inputs={input_A: [[0, 0], [0, 1], [1, 0], [1, 1]],
                                                       input_B: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)

        # Autodiff: SEPARATE INPUT NODES TO SEPARATE OUTPUT NODES WITH ONE OUTPUTPORT EACH
        sizes = {OUTPUT_A: 2, OUTPUT_B: 3}
        nodes = nodes_for_testing_nested_comps(sizes)
        input_A = nodes[INPUT_A]
        input_B = nodes[INPUT_B]
        output_A = nodes[OUTPUT_A]
        output_B = nodes[OUTPUT_B]
        autodiff = AutodiffComposition(pathways=[[input_A, output_A],
                                                 [input_B, output_B]],
                                       name='autodiff')
        result_autodiff_nodes = autodiff.learn(inputs={input_A: [[0, 0], [0, 1], [1, 0], [1, 1]],
                                                       input_B: [[1, 2], [1, 2], [1, 2], [1, 2]]},
                                               learning_rate=.01,
                                               epochs=3,
                                               execution_mode=pnl.ExecutionMode.PyTorch)
        for port, node in zip(result_autodiff_ports, result_autodiff_nodes):
            np.testing.assert_allclose(port, node)


@pytest.mark.pytorch
@pytest.mark.acmisc
@pytest.mark.composition
class TestMiscTrainingFunctionality:

    # test whether pytorch parameters are initialized to be identical to the Autodiff Composition's
    def test_weight_initialization(self):

        # create xor model mechanisms and projections
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection(matrix=np.random.rand(2,10))
        out_map = MappingProjection(matrix=np.random.rand(10,1))

        # put the mechanisms and projections together in an autodiff composition (AC)
        xor = AutodiffComposition()

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # mini version of xor.execute just to build up pytorch representation
        xor._analyze_graph()
        # must manually initialize because _build_pytorch_representation
        # is an internal method
        xor._initialize_from_context(Context(execution_id=xor.default_execution_id))
        xor._build_pytorch_representation(context=xor.default_execution_id)
        # check whether pytorch parameters are identical to projections
        np.testing.assert_allclose(hid_map.parameters.matrix.get(None),
                           list(xor.parameters.pytorch_representation.get(xor).parameters())[0].detach().numpy())
        np.testing.assert_allclose(out_map.parameters.matrix.get(None),
                           list(xor.parameters.pytorch_representation.get(xor).parameters())[1].detach().numpy())

    # test whether processing doesn't interfere with pytorch parameters after training
    def test_training_then_processing(self, autodiff_mode):
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection()
        out_map = MappingProjection()

        xor = AutodiffComposition()

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        xor_inputs = np.array(  # the inputs we will provide to the model
            [[0, 0],
             [0, 1],
             [1, 0],
             [1, 1]])

        xor_targets = np.array(  # the outputs we wish to see from the model
            [[0],
             [1],
             [1],
             [0]])

        # train model for a few epochs
        # results_before_proc = xor.run(inputs={xor_in:xor_inputs},
        #                               targets={xor_out:xor_targets},
        #                               epochs=10)
        results_before_proc = xor.learn(inputs={"inputs": {xor_in:xor_inputs},
                                                "targets": {xor_out:xor_targets},
                                                "epochs": 10},
                                        execution_mode=autodiff_mode)

        # get weight parameters from pytorch
        # pt_weights_hid_bp = xor.parameters.pytorch_representation.get(xor).params[0].detach().numpy().copy()
        # pt_weights_out_bp = xor.parameters.pytorch_representation.get(xor).params[1].detach().numpy().copy()

        #KAM temporarily removed -- will reimplement when pytorch weights can be used in pure PNL execution
        # do processing on a few inputs
        # results_proc = xor.run(inputs={xor_in:xor_inputs})
        # results_proc = xor.run(inputs={"inputs": {xor_in:xor_inputs}})
        #
        # # get weight parameters from pytorch
        # pt_weights_hid_ap = xor.parameters.pytorch_representation.get(xor).params[0].detach().numpy().copy()
        # pt_weights_out_ap = xor.parameters.pytorch_representation.get(xor).params[1].detach().numpy().copy()
        #
        # # check that weight parameters before and after processing are the same
        # np.testing.assert_allclose(pt_weights_hid_bp, pt_weights_hid_ap)
        # np.testing.assert_allclose(pt_weights_out_bp, pt_weights_out_ap)

    @pytest.mark.parametrize(
        'loss, expected', [
            (Loss.CROSS_ENTROPY, [[[0.99330715]], [[0.99933202]], [[0.99933202]], [[0.99985049]]]),
            pytest.param(Loss.L1, [[[0.99330641]], [[0.9993319 ]], [[0.9993319 ]], [[0.99985045]]], marks=pytest.mark.llvm_not_implemented),
            (Loss.MSE, [[[0.99330509]], [[0.99933169]], [[0.99933169]], [[0.9998504]]]),
            pytest.param(Loss.POISSON_NLL, [[[0.99330385]], [[0.99933149]], [[0.99933149]], [[0.99985034]]], marks=pytest.mark.llvm_not_implemented),
        ]
    )
    def test_loss_specs(self, loss, expected, autodiff_mode):
        xor_in = TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        hid_map = MappingProjection()
        out_map = MappingProjection()

        xor = AutodiffComposition(loss_spec=loss)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
        xor_targets = np.array([[0], [1], [1], [0]])

        xor.learn(inputs={"inputs": {xor_in: xor_inputs}, "targets": {xor_out: xor_targets}, "epochs": 10},
                  execution_mode=autodiff_mode)

        tol = {'atol': 2e-6, 'rtol': 2e-6} if loss == Loss.CROSS_ENTROPY else {}
        np.testing.assert_allclose(xor.learning_results, expected, **tol)

    @pytest.mark.llvm_not_implemented
    def test_pytorch_loss_spec(self, autodiff_mode):
        import torch
        ls = torch.nn.SoftMarginLoss(reduction='sum')

        xor_in = TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        hid_map = MappingProjection()
        out_map = MappingProjection()

        xor = AutodiffComposition(loss_spec=ls)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
        xor_targets = np.array([[0], [1], [1], [0]])

        xor.learn(inputs={"inputs": {xor_in: xor_inputs}, "targets": {xor_out: xor_targets}, "epochs": 10},
                  execution_mode=autodiff_mode)
        xor.learn(inputs={"inputs": {xor_in: xor_inputs}, "targets": {xor_out: xor_targets}, "epochs": 10},
                  execution_mode=autodiff_mode)

    @pytest.mark.benchmark(group="Optimizer specs")
    @pytest.mark.parametrize(
        'learning_rate, weight_decay, optimizer_type, expected', [
            (10, 0, 'sgd', [[[0.9863038667851067]], [[0.9944287263151904]], [[0.9934801466163382]], [[0.9979153035411085]]]),
            (1.5, 1, 'sgd', [[[0.33226742]], [[0.4492334]], [[0.75459534]], [[0.44477028]]]),
            (1.5, 1, 'adam', [[[0.43109927]], [[0.33088828]], [[0.40094236]], [[0.57104689]]]),
        ]
    )
    def test_optimizer_specs(self, learning_rate, weight_decay, optimizer_type, expected, autodiff_mode, benchmark):
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection()
        out_map = MappingProjection()

        xor = AutodiffComposition(learning_rate=learning_rate,
                                  optimizer_type=optimizer_type,
                                  weight_decay=weight_decay)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # train model for a few epochs
        results = benchmark(_single_learn_results,
                            xor,
                            inputs={"inputs": {xor_in: [[0, 0], [0, 1], [1, 0], [1, 1]]},
                                    "targets": {xor_out: [[0], [1], [1], [0]]},
                                    "epochs": 10},
                            execution_mode=autodiff_mode)

        # fp32 results are different due to rounding
        if pytest.helpers.llvm_current_fp_precision() == 'fp32' and \
           autodiff_mode is not pnl.ExecutionMode.PyTorch and \
           optimizer_type == 'sgd' and \
           learning_rate == 10:
            expected = [[[0.9918830394744873]], [[0.9982172846794128]], [[0.9978305697441101]], [[0.9994590878486633]]]

        # FIXME: LLVM version is broken with learning rate == 1.5
        if learning_rate != 1.5 or autodiff_mode is pnl.ExecutionMode.PyTorch:
            np.testing.assert_allclose(results, expected)


    # Data for test_learning_rate_assignments():
    default_expected = [[1.07514814, 0.2338579, 1.3005379, 1.05193546, 0.60906245]]      # lr = .0001
    proj_expected = [[1.01985553, 0.78611892, 1.08238162, 1.013176, 0.89058437]]
    constructor_expected = [[1.0546317, 0.43073145, 1.22167948, 1.03716997, 0.70922113]]
    learn_method_expected = [[1.01678782, 0.82604229, 1.06786732, 1.01146838, 0.91116669]]
    no_learning_expected = [[1.07518731, 0.2334628, 1.30069345, 1.05196273, 0.60886095]]
    override_false_expected = [[1.0329353930306493, 0.4250751763513295, 1.1938294699940157,
                                1.0114375211126878, 0.7007802232778438]]

    default_lr = .001

    test_specs_for_learning_rates = [
        #                  |      specify        | cnstr spec | learn() spec |  proj spec  |
        #     condition    | constrctr | learn() | inp   outp |  inp   outp  |  inp   outp |     expected
        #                  |   learning_rate     | projection |  projection  |  projection |     results
        #
        # Test that AutodiffComposition.learning_rate is used
        # ('all_defaults',       False,   False,   None,  None,   None,  None,  None,  None,    default_expected),
        # Test that Projection.learning_rate is used
        ('proj_only',          False,   False,   None,  None,   None,  None,  2.7,   1.4,     proj_expected),
        ('proj_nones',         False,   False,   2.9,    .5,   .66,     1.5,  None, None,     default_expected),
        # Test that Projection.learning_rate == False prevents learning
        ('proj_False_w_nones', False,   False,   None,  None,   None,  None, False, False,    no_learning_expected),
        ('proj_False_w_vals',  False,   False,   2.9,    .5,   .66,     1.5, False, False,    no_learning_expected),
        # Test that Projection.learning_rate is used
        ('proj_like_constr',   False,   False,   None,  None,   None,  None,  2.9,  .5,       constructor_expected),
        # Test that Projection.learning_rate spec supercedes specification in Autodiff constructor
        # ('constr_over_proj',   True,   False,    2.9,    .5,    None,  None,   2.9,  .5,      constructor_expected),
        # Test specs in constructor (with Projecton.learning_rate None or True)
        ('learn_over_proj',    False,   True,    None,  None,   .66,   1.5,   2.7,  1.4,      learn_method_expected),
        # Test that Autodiff constructor spec superceded Projection.learning_rate spec
        ('learn_over_all',     True,   True,     2.9,    .5,    .66,   1.5,   2.7,  1.4,      learn_method_expected),
        # Test that Autodiff constructor spec superceded Projection.learning_rate spec
        ('constructor_only',   True,    False,   2.9,    .5,    None,  None,  None,  True,    constructor_expected),
        # Test specs in learn() method (with Projecton.learning_rate None or True)
        ('learn_method_only',  False,   True,    None,  None,   .66,   1.5,   True,  None,    learn_method_expected),
        # Test that learn() method params supercede constructor params
        ('both',               True,    True,    2.9,    .5,    .66,   1.5,   None,  None,    learn_method_expected),
        # Test that AutodiffComposition.learning_rate is used
        ('constructor_True',   True,    False,   True,  True,   None,  None,  None,  None,    default_expected),
        # Test that AutodiffComposition.learning_rate is used
        ('learn_method_True',  False,   True,    None,  None,   True,  True,  None,  None,    default_expected),
        # Test that AutodiffComposition.learning_rate is used
        ('constructor_None',   True,    False,   None,  None,   None,  None,  None,  None,    default_expected),
        # Test that AutodiffComposition.learning_rate is used
        ('learn_method_None',  False,   True,    None,  None,   None,  None,  None,  None,    default_expected),
        # Test that no learning occurs
        ('constructor_False',  True,    False,   False, False,  None,  None,  None,  None,    no_learning_expected),
        # Test that no learning occurs
        ('learn_method_False', False,   True,    None,  None,   False, False, None,  None,    no_learning_expected),
        # Test that learning does *not* occur
        # ('cnstr_val_lrn_False',True,    True,     2.9,   .5,    False, False, None,  None,    no_learning_expected),
        # Test warning for non-learnable Projection to, within & fron nested Composition
        ('projs_not_lrnable',  True,    True,     2.9,   .5,     .66,   1.5,  None,  None,          None),
        # Test error for assigning False in constructor but learning_rate in learn()
        ('cnstr_False_lrn_val',True,    True,    False, False,   .66,   1.5,  None,  None,    override_false_expected),
        # Test error bad Projection specification
        ('bad_proj',           True,    True,     2.9,   .5,     .66,   1.5,  None,  None,          None),
        # Test error bad learning_rate spec
        ('bad_lr',             True,    True,     2.9,   .5,     .66,   1.5,  None,  None,          None),
    ]
    @pytest.mark.parametrize("condition, use_constructor, use_learn_method, "
                             "constructor_lr_input_proj, constructor_lr_output_proj,"
                             "learn_method_lr_input_proj, learn_method_lr_output_proj,"
                             "projection_lr_input_proj, projection_lr_output_proj,"
                             "expected", test_specs_for_learning_rates,
                             ids=[f"{x[0]}" for x in test_specs_for_learning_rates])
    def test_learning_rate_assignments(self, condition,
                                   use_constructor, use_learn_method,
                                   constructor_lr_input_proj, constructor_lr_output_proj,
                                   learn_method_lr_input_proj, learn_method_lr_output_proj,
                                   projection_lr_input_proj, projection_lr_output_proj,
                                   expected):
        nested_hidden_mech_1 = pnl.ProcessingMechanism(function=pnl.Linear, input_shapes=4, name='nested_1')
        nested_hidden_mech_2 = pnl.ProcessingMechanism(function=Logistic, input_shapes=4, name='nested_2')
        hidden_proj = pnl.MappingProjection(nested_hidden_mech_1, nested_hidden_mech_2, matrix=pnl.IDENTITY_MATRIX)
        nested_comp = pnl.AutodiffComposition([nested_hidden_mech_1, hidden_proj, nested_hidden_mech_2], name="NESTED")
        input_mech = pnl.ProcessingMechanism(input_shapes=3, name='input_mech')
        output_mech = pnl.ProcessingMechanism(input_shapes=5, name='output_mech')
        input_proj = pnl.MappingProjection(name='INPUT PROJECTION',
                                           sender=input_mech,
                                           receiver=nested_hidden_mech_1,
                                           matrix=pnl.RANDOM_CONNECTIVITY_MATRIX,
                                           learning_rate=projection_lr_input_proj)
        output_proj = pnl.MappingProjection(name='OUTPUT PROJECTION',
                                            sender=nested_hidden_mech_2,
                                            receiver=output_mech,
                                            matrix=pnl.RANDOM_CONNECTIVITY_MATRIX,
                                            learning_rate=projection_lr_output_proj)
        inputs={input_mech: [[.1, .2, .3]]}
        targets={output_mech: [[1,1,1,1,1]]}
        constructor_learning_rate_dict = {input_proj: constructor_lr_input_proj,
                                          output_proj: constructor_lr_output_proj}
        learning_method_learning_rate_dict = {input_proj: learn_method_lr_input_proj,
                                              output_proj: learn_method_lr_output_proj}

        opt_params = None

        if condition in {'bad_proj', 'bad_lr'}:
            err_msg = None
            if condition == 'bad_proj':
                opt_params = {condition: .66}
                err_msg = ("The following Projection specified in the 'learning_rate' arg of the learn() method "
                           "for 'OUTER' is not in that Composition or any nested within it: 'bad_proj'.")
            elif condition == 'bad_lr':
                opt_params = {input_proj: condition}
                err_msg = ("A value ('bad_lr') specified in the 'learning_rate' arg of the learn() method for 'OUTER' "
                           "is not valid; it must be an int, float, bool or None.")
            with pytest.raises(AutodiffCompositionError) as error_text:
                outer_comp = pnl.AutodiffComposition(
                    [input_mech, input_proj, nested_comp, output_proj, output_mech], name='OUTER')
                outer_comp.learn(inputs=inputs, targets=targets, learning_rate=opt_params)
            assert err_msg in str(error_text.value)
            return

        elif condition == 'projs_not_lrnable':
            input_proj.learnable = True
            hidden_proj.learnable = False
            output_proj.learnable = False
            opt_params = {input_proj: 1.16, hidden_proj: 21.6, output_proj: 3.99}
            with pytest.raises(AutodiffCompositionError) as error_text:  # Warn, since default_input is NOT set
                outer_comp = pnl.AutodiffComposition([input_mech, input_proj, nested_comp, output_proj, output_mech])
                outer_comp.learn(inputs=inputs, targets=targets, learning_rate=opt_params)
            assert ("Projection ('MappingProjection from nested_1[OutputPort-0] to nested_2[InputPort-0]') specified "
                    "in the dict for the 'learning_rate' arg of the learn() method for 'autodiff_composition' "
                    "is not learnable; check that its 'learnable' attribute is set to 'True' and its learning_rate "
                    "is not 'False', or remove it from the dict."
                    in str(error_text.value))
            return

        elif condition == 'cnstr_False_lrn_val':
            # Validate the learn() method specifications override constructor projection-specific assignments of False
            opt_params = {input_proj: False, hidden_proj: 21.6, output_proj: False}
            outer_comp = pnl.AutodiffComposition([input_mech, input_proj, nested_comp, output_proj, output_mech],
                                                 learning_rate=opt_params)
            opt_params = {input_proj: 0.66, hidden_proj: 21.6, output_proj: 0.5}
            results = outer_comp.learn(inputs=inputs, targets=targets, learning_rate=opt_params, num_trials=2)

            outer_comp_pytorch_rep = outer_comp.parameters.pytorch_representation.get('autodiff_composition')

            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == opt_params[input_proj]
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(hidden_proj) == opt_params[hidden_proj]
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == opt_params[output_proj]
            np.testing.assert_allclose(expected, results)
            return

        # Construct outer Composition
        outer_comp = pnl.AutodiffComposition(
            [input_mech, input_proj, nested_comp, output_proj, output_mech],
            learning_rate=constructor_learning_rate_dict if use_constructor else None,
            name="OUTER")

        # Test after learning with assigned values
        results = outer_comp.learn(
            inputs=inputs, targets=targets,
            learning_rate=learning_method_learning_rate_dict if use_learn_method else None,
            num_trials=2)

        if use_learn_method:
            input_proj_lr = (learn_method_lr_input_proj
                             if ((learn_method_lr_input_proj and learn_method_lr_input_proj is not True)
                                 or learn_method_lr_input_proj is False)
                             else (constructor_lr_input_proj
                                   if (constructor_lr_input_proj and constructor_lr_input_proj is not True)
                                   else (projection_lr_input_proj or self.default_lr)))
            hidden_proj_lr = self.default_lr
            output_proj_lr = (learn_method_lr_output_proj
                              if ((learn_method_lr_output_proj and learn_method_lr_output_proj is not True)
                                  or learn_method_lr_output_proj is False)
                              else (constructor_lr_output_proj
                                    if (constructor_lr_output_proj and constructor_lr_output_proj is not True)
                                    else (constructor_lr_output_proj or self.default_lr)))
        elif use_constructor:
            input_proj_lr = (constructor_lr_input_proj
                             if ((constructor_lr_input_proj and constructor_lr_input_proj is not True)
                                 or constructor_lr_input_proj is False)
                             else (projection_lr_input_proj or self.default_lr))
            hidden_proj_lr = self.default_lr
            output_proj_lr = (constructor_lr_output_proj
                              if ((constructor_lr_output_proj and constructor_lr_output_proj is not True)
                                  or constructor_lr_output_proj is False)
                              else (projection_lr_output_proj or self.default_lr))
        else:
            input_proj_lr = projection_lr_input_proj if projection_lr_input_proj is not None else self.default_lr
            hidden_proj_lr = self.default_lr
            output_proj_lr = projection_lr_output_proj if projection_lr_output_proj is not None else self.default_lr

        outer_comp_pytorch_rep = outer_comp.parameters.pytorch_representation.get('OUTER')
        assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == input_proj_lr
        assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(hidden_proj) == hidden_proj_lr
        assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == output_proj_lr
        np.testing.assert_allclose(expected, results)

        # Learning rate should return to default values if not specified again
        outer_comp.learn(inputs=inputs, targets=targets)
        if condition == 'both':
            # Should return to defaults specified in constructor (even though specified in previous call to learning)
            assert len(outer_comp_pytorch_rep.optimizer.param_groups) == 3
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == constructor_lr_input_proj or self.default_lr
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(hidden_proj) == 0.001
            assert (outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == constructor_lr_output_proj or self.default_lr)
        elif condition == 'learn_method':
            # Should return to default for optimizer (since none specified for constructor)
            assert len(outer_comp_pytorch_rep.optimizer.param_groups) == 3
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(input_proj) == self.default_lr
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(hidden_proj) == self.default_lr
            assert outer_comp_pytorch_rep.get_torch_learning_rate_for_projection(output_proj) == self.default_lr


    default = .001
    test_specs_for_learning_rate_inheritance = [
        #                                                                  NOTE: the specs below are for values
        #                                                                        expected after construction; values
        #                                                                        after learn() are handled in the test
        #  condition    p_1_lr  p_2_lr  pathway_lr  in_cmp_lr  out_cmp_lr  out_lrn_lr  exp_p_1_in exp_p2_in
        ('defaults',       None,   None,     None,      None,      None,  NotImplemented, default,  default),
        # projection-specific specs takes precedence if no learn() method specs
        ('proj_lr_nimp',  1.414,     7,      None,      6.02,       2.7,  NotImplemented,  1.414,       7),
        ('proj_lr_none',  1.414,     7,      None,      6.02,       2.7,      None,        1.414,       7),
        # projection-specific specs takes precedence over pathway, but pathay takes precedence over comp lr's
        ('pathway_lr',    1.414,   None,     2.99,      6.02,       1.6,      3.14,        1.414,      2.99),
        # learn() method takes precedence, and specifying None for Projections forces them to use relevant default
        #        NOTE:  out_lrn_lr only applied to inner_proj_1 or inner_proj_2
        ('learn_only',     None,   None,     None,      None,      None,      3.14,       default,  default),
        ('inr_p2_none',   1.414,   None,     None,      None,      None,      3.14,        1.414,   default),
        ('inr_outr',      1.414,   None,     None,      6.02,      None,      3.14,        1.414,     6.02),
        ('inr_none_nimp',  None,     7,      None,      6.02,      None,  NotImplemented,   6.02,       7),
        ('inr_none_none',  None,     7,      None,      6.02,      None,      None,         6.02,       7),
        ('inr_lr_nimp',   1.414,     7,      None,      6.02,      None,  NotImplemented,  1.414,       7),
        ('inr_lr_none',   1.414,     7,      None,      6.02,      None,      None,        1.414,       7),
        ('outr_comp_lr',   None,     7,      None,      6.02,       2.7,      None,         6.02,       7),
        ('outr_comp_lr',   None,     7,      None,      None,       2.7,      None,       default,      7),
    ]
    @pytest.mark.parametrize("_condition, proj_1_lr, proj_2_lr, pathway_lr, "
                             "inner_comp_lr, outer_comp_lr, outer_learn_lr, "
                             "expected_proj_1_inner, expected_proj_2_inner",
                             test_specs_for_learning_rate_inheritance,
                             ids=[f"{x[0]}" for x in test_specs_for_learning_rate_inheritance])
    def test_learning_rate_inheritance_in_nested_comp(self, _condition, proj_1_lr, proj_2_lr, pathway_lr,
                                                      inner_comp_lr, outer_comp_lr, outer_learn_lr,
                                                      expected_proj_1_inner, expected_proj_2_inner):

        # Construct inner Composition
        inner_node_input = pnl.ProcessingMechanism(name="INPUT NODE")
        inner_node_hidden = pnl.ProcessingMechanism(name="HIDDEN NODE")
        inner_node_output = pnl.ProcessingMechanism(name="OUTPUT NODE")
        inner_proj_1 = pnl.MappingProjection(inner_node_input,
                                           inner_node_hidden,
                                           learning_rate = proj_1_lr,
                                           name="INNER PROJECTION 1")
        inner_proj_2 = pnl.MappingProjection(inner_node_hidden,
                                           inner_node_output,
                                           learning_rate = proj_2_lr,
                                           name="INNER PROJECTION 2")
        # inner_pathway = [inner_node_input, inner_proj_1, inner_node_hidden, inner_proj_2, inner_node_output]
        # inner_comp = AutodiffComposition(pathways=inner_pathway,
        #                                  learning_rate=inner_comp_lr,
        #                                  name="INNER COMP")

        # BREADCRUMB
        inner_comp = AutodiffComposition(learning_rate=inner_comp_lr,
                                         name="INNER COMP")
        inner_comp.add_backpropagation_learning_pathway([inner_node_input,
                                                         inner_proj_1,
                                                         inner_node_hidden,
                                                         inner_proj_2,
                                                         inner_node_output],
                                                        learning_rate = pathway_lr)
        # BREADCRUMB
        # Check inner_comp assignments from constructor
        inner_pytorch_rep = inner_comp._build_pytorch_representation()
        # Ensure that params were assigned appropriate lr for the inner_comp
        assert inner_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_1) == expected_proj_1_inner
        assert inner_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_2) == expected_proj_2_inner

        # Construct outer Composition with nested inner
        outer_node = pnl.ProcessingMechanism(name="OUTER NODE")

        # BREADCRUMB
        outer_comp = AutodiffComposition([inner_comp, outer_node], learning_rate=outer_comp_lr, name="OUTER COMP")

        # Check outer_comp assignments from constructor
        # BREADCRUMB
        outer_pytorch_rep = outer_comp._build_pytorch_representation()
        outer_proj = outer_comp.nodes[-1].afferents[0]
        assert outer_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_1) == expected_proj_1_inner
        assert outer_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_2) == expected_proj_2_inner
        assert outer_pytorch_rep.get_torch_learning_rate_for_projection(outer_proj) == outer_comp_lr or self.default

        # Check outer_comp assignments after learn() method
        if outer_learn_lr == NotImplemented:
            learning_rate_arg = None
            # proj_1_expected = proj_1_lr or inner_comp_lr or outer_comp_lr or self.default
            proj_1_expected = proj_1_lr or inner_comp_lr or outer_comp_lr or self.default
            proj_2_expected = proj_2_lr or inner_comp_lr or outer_comp_lr or self.default
            outer_proj_expected = outer_comp_lr or self.default
        elif outer_learn_lr is None:
            learning_rate_arg = {inner_proj_1: None,
                             outer_proj: None}
            proj_1_expected = inner_comp_lr or proj_1_lr or outer_comp_lr
            proj_2_expected = inner_comp_lr or proj_2_lr or self.default
            outer_proj_expected = outer_comp_lr or self.default
        else:
            learning_rate_arg = {inner_proj_1: outer_learn_lr,
                             outer_proj: outer_learn_lr}
            proj_1_expected = outer_learn_lr or proj_1_lr or self.default
            proj_2_expected = outer_learn_lr or proj_2_lr or self.default
            outer_proj_expected = outer_learn_lr or outer_learn_lr or self.default

        # BREADCRUMB
        outer_comp.learn(inputs={inner_node_input:[[1]]},
                         learning_rate=learning_rate_arg)
        learn_pytorch_rep = outer_comp.parameters.pytorch_representation.get('OUTER COMP')
        assert learn_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_1) == proj_1_expected
        assert learn_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_2) == expected_proj_2_inner
        assert learn_pytorch_rep.get_torch_learning_rate_for_projection(outer_proj) == outer_proj_expected
        # Check that _learing_params_for_execution works properly
        assert learn_pytorch_rep._torch_params_for_execution[inner_proj_1] == proj_1_expected
        assert learn_pytorch_rep._torch_params_for_execution[inner_proj_2] == expected_proj_2_inner
        assert learn_pytorch_rep._torch_params_for_execution[outer_proj] == outer_proj_expected

        # BREADCRUMB
        # Check that learning_rates return to those at construction after another call to learn() w/o learning_rate_arg
        outer_comp.learn(inputs={inner_node_input:[[1]]})
        assert learn_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_1) == expected_proj_1_inner
        assert learn_pytorch_rep.get_torch_learning_rate_for_projection(inner_proj_2) == expected_proj_2_inner
        assert learn_pytorch_rep.get_torch_learning_rate_for_projection(outer_proj) == outer_comp_lr or self.default

    def test_pytorch_identicality_of_learning_rates_nested(self):
        """Test ideinticality with assignment of learning rates in learning_rate of nested Composition"""
        import torch
        entry_torch_dtype = torch.get_default_dtype()
        torch.set_default_dtype(torch.float64)

        INPUT_SIZE = 3
        HIDDEN_SIZE = 2
        NESTED_1_SIZE = 3
        NESTED_2_SIZE = 2
        OUTPUT_SIZE = 5
        LEARNING_RATE = .001
        INPUT_LEARNING_RATE = 1.414
        HIDDEN_LEARNING_RATE = 6.02
        OUTPUT_LEARNING_RATE = 2.7
        input_stim = [[1.,2.,3.]]
        target_stim = [[1.,1.,1.,1.,1.]]

        #  TorchModel:
        #  input_wts      hidden_wts       output_wts
        #    ->     input     ->    hidden    ->    output
        #        INPUT_SIZE       HIDDEN_SIZE     OUTPUT_SIZE
        #
        #  PNL Model:
        #       input_proj          nested_proj          output_proj
        #  input    ->    [ nested_1    ->     nested_2 ]   ->     output
        # INPUT_SIZE       INPUT_SIZE        HIDDEN_SIZE         OUTPUT_SIZE
        #
        # Mapping of weights between models:
        #   Torch           PNL
        # input_wts   ->   input_proj
        # hidden_wts  ->   nested_proj
        # output_wts  ->   output_proj

        # Set up and run torch model ----------------------------------------------------------------------

        # Set up torch model
        class TorchModel(torch.nn.Module):
            def __init__(self):
                super(TorchModel, self).__init__()
                self.input = torch.nn.Linear(INPUT_SIZE, INPUT_SIZE, bias=False)
                self.hidden = torch.nn.Linear(INPUT_SIZE, HIDDEN_SIZE, bias=False)
                self.output = torch.nn.Linear(HIDDEN_SIZE, OUTPUT_SIZE, bias=False)

            def forward(self, x):
                after_input = self.input(x)
                after_hidden = self.hidden(after_input)
                after_output = self.output(after_hidden)
                return after_output

        # Set up Torch optimizer
        torch_model = TorchModel()
        torch_optimizer = torch.optim.SGD(lr=LEARNING_RATE, params=torch_model.parameters())
        del torch_optimizer.param_groups[0]
        torch_param_names_map = {p[0]:p[1] for p in torch_model.named_parameters()}
        torch_optimizer.add_param_group({'params': torch_param_names_map['input.weight'], 'lr': INPUT_LEARNING_RATE})
        torch_optimizer.add_param_group({'params': torch_param_names_map['hidden.weight'], 'lr': HIDDEN_LEARNING_RATE})
        torch_optimizer.add_param_group({'params': torch_param_names_map['output.weight'], 'lr': OUTPUT_LEARNING_RATE})
        loss_fct = torch.nn.MSELoss(reduction='mean')

        # Get initial weights (to initialize autodiff below with same initial conditions)
        torch_input_initial_weights = torch_model.state_dict()['input.weight'].T.detach().cpu().numpy().copy()
        torch_hidden_initial_weights = torch_model.state_dict()['hidden.weight'].T.detach().cpu().numpy().copy()
        torch_output_initial_weights = torch_model.state_dict()['output.weight'].T.detach().cpu().numpy().copy()

        # Execute model without learning
        torch_result_before_learning = torch_model(torch.tensor(np.array(input_stim)))

        # Compute loss and update weights
        torch_optimizer.zero_grad()
        torch_loss = loss_fct(torch_result_before_learning, torch.tensor(target_stim))
        torch_loss.backward()
        torch_optimizer.step()

        # Get output after learning
        torch_result_after_learning = torch_model(torch.tensor(np.array(input_stim)))

        # Set up and run PNL Autodiff model ------------------------------------------------------------

        # Nested Composition
        nested_hidden_mech_1 = pnl.ProcessingMechanism(input_shapes=INPUT_SIZE, name='nested_1')
        nested_hidden_mech_2 = pnl.ProcessingMechanism(input_shapes=HIDDEN_SIZE, name='nested_2')
        nested_proj = pnl.MappingProjection(nested_hidden_mech_1, nested_hidden_mech_2,
                                            matrix=pnl.RANDOM_CONNECTIVITY_MATRIX)
        nested_comp = pnl.AutodiffComposition([nested_hidden_mech_1, nested_proj, nested_hidden_mech_2],
                                              name='NESTED COMP')
        # Outer Composition
        input_mech = pnl.ProcessingMechanism(input_shapes=INPUT_SIZE, name='input_mech')
        output_mech = pnl.ProcessingMechanism(input_shapes=OUTPUT_SIZE, name='output_mech')
        input_proj = pnl.MappingProjection(input_mech, nested_hidden_mech_1, matrix=pnl.RANDOM_CONNECTIVITY_MATRIX)
        output_proj = pnl.MappingProjection(nested_hidden_mech_2, output_mech, matrix=pnl.RANDOM_CONNECTIVITY_MATRIX)
        inputs={input_mech: input_stim}
        targets={output_mech: target_stim}
        learning_rate={pnl.DEFAULT_LEARNING_RATE: LEARNING_RATE,
                       input_proj: INPUT_LEARNING_RATE,
                       nested_proj: HIDDEN_LEARNING_RATE,
                       output_proj: OUTPUT_LEARNING_RATE}
        outer_comp = pnl.AutodiffComposition(
            name='OUTER COMP',
            pathways=[input_mech, input_proj, nested_comp, output_proj, output_mech],
            learning_rate=learning_rate
        )

        # Assign initial weights from TorchModel to PNL:
        #   Torch           PNL
        # input_wts   ->   input_proj
        # hidden_wts  ->   nested_proj
        # output_wts  ->   output_proj
        outer_comp.set_weights(outer_comp.projections[0], torch_input_initial_weights)
        nested_comp.set_weights(nested_proj, torch_hidden_initial_weights)
        outer_comp.set_weights(outer_comp.projections[1], torch_output_initial_weights)
        target_mechs = outer_comp.infer_backpropagation_learning_pathways(pnl.ExecutionMode.PyTorch)

        # Execute autodiff with learning (but not weight updates yet)
        autodiff_result_before_learning = outer_comp.learn(inputs={input_mech:input_stim, target_mechs[0]: target_stim},
                                                           execution_mode=pnl.ExecutionMode.PyTorch,
                                                           )
        # Get results after learning (with weight updates)
        autodiff_result_after_learning = outer_comp.run(inputs={input_mech:input_stim},
                                                        execution_mode=pnl.ExecutionMode.PyTorch)

        # Compare results ------------------------------------------------------------

        # Test of forward pass (without effects of learning yet):
        np.testing.assert_allclose(torch_result_before_learning.detach().numpy(),
                                   autodiff_result_before_learning, atol=1e-8)

        # Test of execution after backward pass (learning):
        np.testing.assert_allclose(torch_loss.detach().numpy(), outer_comp.torch_losses.squeeze())
        np.testing.assert_allclose(torch_result_after_learning.detach().numpy(),
                                   autodiff_result_after_learning, atol=1e-8)

        torch.set_default_dtype(entry_torch_dtype)

    # test whether pytorch parameters and projections are kept separate (at diff. places in memory)
    def test_params_stay_separate(self, autodiff_mode):
        if autodiff_mode is not pnl.ExecutionMode.PyTorch:
            pytest.skip("Compiled weights are always copied back!")

        xor_in = TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        hid_m = np.random.rand(2,10)
        out_m = np.random.rand(10,1)

        hid_map = MappingProjection(name='hid_map',
                                    matrix=hid_m.copy(),
                                    sender=xor_in,
                                    receiver=xor_hid)

        out_map = MappingProjection(name='out_map',
                                    matrix=out_m.copy(),
                                    sender=xor_hid,
                                    receiver=xor_out)

        xor = AutodiffComposition(learning_rate=10.0,
                                  optimizer_type="sgd")

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        xor_inputs = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])

        xor_targets = np.array([[0], [1], [1], [0]])

        # train the model for a few epochs
        xor.learn(inputs={"inputs": {xor_in:xor_inputs}, "targets": {xor_out:xor_targets},"epochs": 10},
                  execution_mode=autodiff_mode)

        # get weight parameters from pytorch
        # pt_weights_hid = xor.parameters.pytorch_representation.get(xor).params[0].detach().numpy().copy()
        # pt_weights_out = xor.parameters.pytorch_representation.get(xor).params[1].detach().numpy().copy()
        pt_weights_hid = list(xor.parameters.pytorch_representation.get(xor).parameters())[0].detach().numpy().copy()
        pt_weights_out = list(xor.parameters.pytorch_representation.get(xor).parameters())[1].detach().numpy().copy()
        # assert that projections are still what they were initialized as
        np.testing.assert_allclose(hid_map.parameters.matrix.get(None), hid_m)
        np.testing.assert_allclose(out_map.parameters.matrix.get(None), out_m)

        # assert that projections didn't change during training with the pytorch
        # parameters (they should now be different)
        assert not np.allclose(pt_weights_hid, hid_map.parameters.matrix.get(None))
        assert not np.allclose(pt_weights_out, out_map.parameters.matrix.get(None))


@pytest.mark.pytorch
@pytest.mark.actime
class TestTrainingTime:

    @pytest.mark.skip
    @pytest.mark.parametrize(
        'eps, opt', [
            (1, 'sgd'),
            (10, 'sgd'),
            (100, 'sgd')
        ]
    )
    def test_and_training_time(self, eps, opt,autodiff_mode):

        # SET UP MECHANISMS FOR COMPOSITION

        and_in = TransferMechanism(name='and_in',
                                   default_variable=np.zeros(2))

        and_out = TransferMechanism(name='and_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        # SET UP MECHANISMS FOR SYSTEM

        and_in_sys = TransferMechanism(name='and_in_sys',
                                       default_variable=np.zeros(2))

        and_out_sys = TransferMechanism(name='and_out_sys',
                                        default_variable=np.zeros(1),
                                        function=Logistic())

        # SET UP PROJECTIONS FOR COMPOSITION

        and_map = MappingProjection(name='and_map',
                                    matrix=np.random.rand(2, 1),
                                    sender=and_in,
                                    receiver=and_out)

        # SET UP PROJECTIONS FOR SYSTEM

        and_map_sys = MappingProjection(name='and_map_sys',
                                        matrix=and_map.matrix.base.copy(),
                                        sender=and_in_sys,
                                        receiver=and_out_sys)

        # SET UP COMPOSITION

        and_net = AutodiffComposition()

        and_net.add_node(and_in)
        and_net.add_node(and_out)

        and_net.add_projection(sender=and_in, projection=and_map, receiver=and_out)

        # SET UP INPUTS AND TARGETS

        and_inputs = np.zeros((4,2))
        and_inputs[0] = [0, 0]
        and_inputs[1] = [0, 1]
        and_inputs[2] = [1, 0]
        and_inputs[3] = [1, 1]

        and_targets = np.zeros((4,1))
        and_targets[0] = [0]
        and_targets[1] = [1]
        and_targets[2] = [1]
        and_targets[3] = [0]

        # TIME TRAINING FOR COMPOSITION

        start = timeit.default_timer()
        result = and_net.run(inputs={and_in:and_inputs},
                             targets={and_out:and_targets},
                             epochs=eps,
                             learning_rate=0.1,
                             controller=opt,
                             execution_mode=autodiff_mode)
        end = timeit.default_timer()
        comp_time = end - start

        msg = 'Training XOR model as AutodiffComposition for {0} epochs took {1} seconds'.format(eps, comp_time)
        print(msg)
        print("\n")
        logger.info(msg)

    @pytest.mark.skip
    @pytest.mark.parametrize(
        'eps, opt', [
            (1, 'sgd'),
            (10, 'sgd'),
            (100, 'sgd')
        ]
    )
    def test_xor_training_time(self, eps, opt, autodiff_mode):

        # SET UP MECHANISMS FOR COMPOSITION

        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        # SET UP MECHANISMS FOR SYSTEM

        xor_in_sys = TransferMechanism(name='xor_in_sys',
                                       default_variable=np.zeros(2))

        xor_hid_sys = TransferMechanism(name='xor_hid_sys',
                                        default_variable=np.zeros(10),
                                        function=Logistic())

        xor_out_sys = TransferMechanism(name='xor_out_sys',
                                        default_variable=np.zeros(1),
                                        function=Logistic())

        # SET UP PROJECTIONS FOR COMPOSITION

        hid_map = MappingProjection(name='hid_map',
                                    matrix=np.random.rand(2,10),
                                    sender=xor_in,
                                    receiver=xor_hid)

        out_map = MappingProjection(name='out_map',
                                    matrix=np.random.rand(10,1),
                                    sender=xor_hid,
                                    receiver=xor_out)

        # SET UP PROJECTIONS FOR SYSTEM

        hid_map_sys = MappingProjection(name='hid_map_sys',
                                        matrix=hid_map.matrix.base.copy(),
                                        sender=xor_in_sys,
                                        receiver=xor_hid_sys)

        out_map_sys = MappingProjection(name='out_map_sys',
                                        matrix=out_map.matrix.base.copy(),
                                        sender=xor_hid_sys,
                                        receiver=xor_out_sys)

        # SET UP COMPOSITION

        xor = AutodiffComposition(execution_mode=autodiff_mode)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # SET UP INPUTS AND TARGETS

        xor_inputs = np.array(  # the inputs we will provide to the model
            [[0, 0],
             [0, 1],
             [1, 0],
             [1, 1]])

        xor_targets = np.array(  # the outputs we wish to see from the model
            [[0],
             [1],
             [1],
             [0]])

        # TIME TRAINING FOR COMPOSITION

        start = timeit.default_timer()
        result = xor.run(inputs={xor_in:xor_inputs},
                         targets={xor_out:xor_targets},
                         epochs=eps,
                         learning_rate=0.1,
                         controller=opt,
                         execution_mode=autodiff_mode)
        end = timeit.default_timer()
        comp_time = end - start

        # SET UP SYSTEM

        # xor_process = Process(pathway=[xor_in_sys,
        #                                hid_map_sys,
        #                                xor_hid_sys,
        #                                out_map_sys,
        #                                xor_out_sys],
        #                       learning=pnl.LEARNING)

        xor_process = Composition(pathways=([xor_in_sys,
                                       hid_map_sys,
                                       xor_hid_sys,
                                       out_map_sys,
                                       xor_out_sys], BackPropagation))

        msg = 'Training XOR model as AutodiffComposition for {eps} epochs took {comp_time} seconds.'
        print(msg)
        print("\n")
        logger.info(msg)

    @pytest.mark.skip
    @pytest.mark.parametrize(
            'eps, opt', [
            (1, 'sgd'),
            (10, 'sgd'),
            (100, 'sgd')
        ]
    )
    def test_semantic_net_training_time(self, eps, opt):

        # SET UP MECHANISMS FOR COMPOSITION:

        nouns_in = TransferMechanism(name="nouns_input",
                                     default_variable=np.zeros(8))

        rels_in = TransferMechanism(name="rels_input",
                                    default_variable=np.zeros(3))

        h1 = TransferMechanism(name="hidden_nouns",
                               default_variable=np.zeros(8),
                               function=Logistic())

        h2 = TransferMechanism(name="hidden_mixed",
                               default_variable=np.zeros(15),
                               function=Logistic())

        out_sig_I = TransferMechanism(name="sig_outs_I",
                                      default_variable=np.zeros(8),
                                      function=Logistic())

        out_sig_is = TransferMechanism(name="sig_outs_is",
                                       default_variable=np.zeros(12),
                                       function=Logistic())

        out_sig_has = TransferMechanism(name="sig_outs_has",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        out_sig_can = TransferMechanism(name="sig_outs_can",
                                        default_variable=np.zeros(9),
                                        function=Logistic())

        # SET UP MECHANISMS FOR SYSTEM

        nouns_in_sys = TransferMechanism(name="nouns_input_sys",
                                         default_variable=np.zeros(8))

        rels_in_sys = TransferMechanism(name="rels_input_sys",
                                        default_variable=np.zeros(3))

        h1_sys = TransferMechanism(name="hidden_nouns_sys",
                                   default_variable=np.zeros(8),
                                   function=Logistic())

        h2_sys = TransferMechanism(name="hidden_mixed_sys",
                                   default_variable=np.zeros(15),
                                   function=Logistic())

        out_sig_I_sys = TransferMechanism(name="sig_outs_I_sys",
                                          default_variable=np.zeros(8),
                                          function=Logistic())

        out_sig_is_sys = TransferMechanism(name="sig_outs_is_sys",
                                           default_variable=np.zeros(12),
                                           function=Logistic())

        out_sig_has_sys = TransferMechanism(name="sig_outs_has_sys",
                                            default_variable=np.zeros(9),
                                            function=Logistic())

        out_sig_can_sys = TransferMechanism(name="sig_outs_can_sys",
                                            default_variable=np.zeros(9),
                                            function=Logistic())

        # SET UP PROJECTIONS FOR COMPOSITION

        map_nouns_h1 = MappingProjection(matrix=np.random.rand(8,8),
                                         name="map_nouns_h1",
                                         sender=nouns_in,
                                         receiver=h1)

        map_rels_h2 = MappingProjection(matrix=np.random.rand(3,15),
                                        name="map_rel_h2",
                                        sender=rels_in,
                                        receiver=h2)

        map_h1_h2 = MappingProjection(matrix=np.random.rand(8,15),
                                      name="map_h1_h2",
                                      sender=h1,
                                      receiver=h2)

        map_h2_I = MappingProjection(matrix=np.random.rand(15,8),
                                     name="map_h2_I",
                                     sender=h2,
                                     receiver=out_sig_I)

        map_h2_is = MappingProjection(matrix=np.random.rand(15,12),
                                      name="map_h2_is",
                                      sender=h2,
                                      receiver=out_sig_is)

        map_h2_has = MappingProjection(matrix=np.random.rand(15,9),
                                       name="map_h2_has",
                                       sender=h2,
                                       receiver=out_sig_has)

        map_h2_can = MappingProjection(matrix=np.random.rand(15,9),
                                       name="map_h2_can",
                                       sender=h2,
                                       receiver=out_sig_can)

        # SET UP PROJECTIONS FOR SYSTEM

        map_nouns_h1_sys = MappingProjection(matrix=map_nouns_h1.matrix.base.copy(),
                                             name="map_nouns_h1_sys",
                                             sender=nouns_in_sys,
                                             receiver=h1_sys)

        map_rels_h2_sys = MappingProjection(matrix=map_rels_h2.matrix.base.copy(),
                                        name="map_relh2_sys",
                                        sender=rels_in_sys,
                                        receiver=h2_sys)

        map_h1_h2_sys = MappingProjection(matrix=map_h1_h2.matrix.base.copy(),
                                          name="map_h1_h2_sys",
                                          sender=h1_sys,
                                          receiver=h2_sys)

        map_h2_I_sys = MappingProjection(matrix=map_h2_I.matrix.base.copy(),
                                         name="map_h2_I_sys",
                                         sender=h2_sys,
                                         receiver=out_sig_I_sys)

        map_h2_is_sys = MappingProjection(matrix=map_h2_is.matrix.base.copy(),
                                          name="map_h2_is_sys",
                                          sender=h2_sys,
                                          receiver=out_sig_is_sys)

        map_h2_has_sys = MappingProjection(matrix=map_h2_has.matrix.base.copy(),
                                           name="map_h2_has_sys",
                                           sender=h2_sys,
                                           receiver=out_sig_has_sys)

        map_h2_can_sys = MappingProjection(matrix=map_h2_can.matrix.base.copy(),
                                           name="map_h2_can_sys",
                                           sender=h2_sys,
                                           receiver=out_sig_can_sys)

        # COMPOSITION FOR SEMANTIC NET

        sem_net = AutodiffComposition()

        sem_net.add_node(nouns_in)
        sem_net.add_node(rels_in)
        sem_net.add_node(h1)
        sem_net.add_node(h2)
        sem_net.add_node(out_sig_I)
        sem_net.add_node(out_sig_is)
        sem_net.add_node(out_sig_has)
        sem_net.add_node(out_sig_can)

        sem_net.add_projection(sender=nouns_in, projection=map_nouns_h1, receiver=h1)
        sem_net.add_projection(sender=rels_in, projection=map_rels_h2, receiver=h2)
        sem_net.add_projection(sender=h1, projection=map_h1_h2, receiver=h2)
        sem_net.add_projection(sender=h2, projection=map_h2_I, receiver=out_sig_I)
        sem_net.add_projection(sender=h2, projection=map_h2_is, receiver=out_sig_is)
        sem_net.add_projection(sender=h2, projection=map_h2_has, receiver=out_sig_has)
        sem_net.add_projection(sender=h2, projection=map_h2_can, receiver=out_sig_can)

        # INPUTS & OUTPUTS FOR SEMANTIC NET:

        nouns = ['oak', 'pine', 'rose', 'daisy', 'canary', 'robin', 'salmon', 'sunfish']
        relations = ['is', 'has', 'can']
        is_list = ['living', 'living thing', 'plant', 'animal', 'tree', 'flower', 'bird', 'fish', 'big', 'green', 'red',
                   'yellow']
        has_list = ['roots', 'leaves', 'bark', 'branches', 'skin', 'feathers', 'wings', 'gills', 'scales']
        can_list = ['grow', 'move', 'swim', 'fly', 'breathe', 'breathe underwater', 'breathe air', 'walk', 'photosynthesize']

        nouns_input = np.identity(len(nouns))

        rels_input = np.identity(len(relations))

        truth_nouns = np.identity(len(nouns))

        truth_is = np.zeros((len(nouns), len(is_list)))

        truth_is[0, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0]
        truth_is[1, :] = [1, 1, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0]
        truth_is[2, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
        truth_is[3, :] = [1, 1, 1, 0, 0, 1, 0, 0, 0, 0, 0, 0]
        truth_is[4, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[5, :] = [1, 1, 0, 1, 0, 0, 1, 0, 0, 0, 0, 1]
        truth_is[6, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 1, 0]
        truth_is[7, :] = [1, 1, 0, 1, 0, 0, 0, 1, 1, 0, 0, 0]

        truth_has = np.zeros((len(nouns), len(has_list)))

        truth_has[0, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[1, :] = [1, 1, 1, 1, 0, 0, 0, 0, 0]
        truth_has[2, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[3, :] = [1, 1, 0, 0, 0, 0, 0, 0, 0]
        truth_has[4, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[5, :] = [0, 0, 0, 0, 1, 1, 1, 0, 0]
        truth_has[6, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]
        truth_has[7, :] = [0, 0, 0, 0, 0, 0, 0, 1, 1]

        truth_can = np.zeros((len(nouns), len(can_list)))

        truth_can[0, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[1, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[2, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[3, :] = [1, 0, 0, 0, 0, 0, 0, 0, 1]
        truth_can[4, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[5, :] = [1, 1, 0, 1, 1, 0, 1, 1, 0]
        truth_can[6, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]
        truth_can[7, :] = [1, 1, 1, 0, 1, 1, 0, 0, 0]

        # SETTING UP DICTIONARIES OF INPUTS/OUTPUTS FOR SEMANTIC NET

        inputs_dict = {}
        inputs_dict[nouns_in] = []
        inputs_dict[rels_in] = []

        targets_dict = {}
        targets_dict[out_sig_I] = []
        targets_dict[out_sig_is] = []
        targets_dict[out_sig_has] = []
        targets_dict[out_sig_can] = []

        for i in range(len(nouns)):
            for j in range(len(relations)):
                inputs_dict[nouns_in].append(nouns_input[i])
                inputs_dict[rels_in].append(rels_input[j])
                targets_dict[out_sig_I].append(truth_nouns[i])
                targets_dict[out_sig_is].append(truth_is[i])
                targets_dict[out_sig_has].append(truth_has[i])
                targets_dict[out_sig_can].append(truth_can[i])

        inputs_dict_sys = {}
        inputs_dict_sys[nouns_in_sys] = inputs_dict[nouns_in]
        inputs_dict_sys[rels_in_sys] = inputs_dict[rels_in]

        targets_dict_sys = {}
        targets_dict_sys[out_sig_I_sys] = targets_dict[out_sig_I]
        targets_dict_sys[out_sig_is_sys] = targets_dict[out_sig_is]
        targets_dict_sys[out_sig_has_sys] = targets_dict[out_sig_has]
        targets_dict_sys[out_sig_can_sys] = targets_dict[out_sig_can]

        # TIME TRAINING FOR COMPOSITION

        start = timeit.default_timer()
        result = sem_net.run(inputs=inputs_dict,
                             targets=targets_dict,
                             epochs=eps,
                             learning_rate=0.1,
                             controller=opt)
        end = timeit.default_timer()
        comp_time = end - start

        msg = 'Training Semantic net as AutodiffComposition for {0} epochs took {1} seconds'.format(eps, comp_time)
        print(msg)
        print("\n")
        logger.info(msg)


@pytest.mark.pytorch
def test_autodiff_saveload(tmp_path):
    def create_xor():
        # create xor model mechanisms and projections
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection(matrix=np.random.rand(2,10), name='hid_map')
        out_map = MappingProjection(matrix=np.random.rand(10,1), name='out_map')

        # put the mechanisms and projections together in an autodiff composition (AC)
        xor = AutodiffComposition()

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)
        return xor

    np.random.seed(0)
    xor1 = create_xor()
    xor1_outputs = xor1.run(inputs=[0,0])

    # save
    # path = xor1.save()
    path = xor1.save(os.path.join(tmp_path, 'xor_1.pnl'))

    # del xor1
    pnl.clear_registry()

    # load
    np.random.seed(1)
    xor2 = create_xor()
    xor2_outputs_pre = xor2.run(inputs=[0,0])
    # xor2.load(os.path.join(tmp_path, 'xor_1.pnl'))
    xor2.load(path)
    xor2_outputs_post = xor2.run(inputs=[0,0])


    # sanity check - make sure xor2 weights differ
    assert not np.allclose(xor2_outputs_pre, xor2_outputs_post, atol=1e-9)

    # make sure loaded model is identical, and used during run
    np.testing.assert_allclose(xor1_outputs, xor2_outputs_post, atol=1e-9)


@pytest.mark.pytorch
@pytest.mark.aclogging
class TestACLogging:

    @pytest.mark.parametrize('minibatch_size', [1, 2])
    def test_autodiff_logging(self, minibatch_size):
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection()
        out_map = MappingProjection()

        xor = AutodiffComposition()

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        xor_in.set_log_conditions('value', pnl.LogCondition.TRIAL)
        xor_hid.set_log_conditions('value', pnl.LogCondition.TRIAL)
        xor_out.set_log_conditions('value', pnl.LogCondition.TRIAL)

        hid_map.set_log_conditions('matrix', pnl.LogCondition.TRIAL)
        out_map.set_log_conditions('matrix', pnl.LogCondition.TRIAL)

        xor_inputs = np.array(  # the inputs we will provide to the model
            [[[0, 0]],
             [[0, 1]],
             [[1, 0]],
             [[1, 1]]])

        xor_targets = np.array(  # the outputs we wish to see from the model
            [[0],
             [1],
             [1],
             [0]])

        # train model for a few epochs
        num_epochs = 10
        xor.learn(inputs={"inputs": {xor_in: xor_inputs},
                          "targets": {xor_out: xor_targets},
                          "epochs": num_epochs},
                  synch_projection_matrices_with_torch=pnl.MINIBATCH,
                  synch_results_with_torch=pnl.MINIBATCH,
                  # synch_results_with_torch=pnl.RUN,
                  execution_mode=pnl.ExecutionMode.PyTorch,
                  minibatch_size=minibatch_size)

        exec_id = xor.default_execution_id

        in_np_dict_vals = xor_in.log.nparray_dictionary()[exec_id]['value']
        in_np_vals = xor_in.log.nparray()[1][1][4][1:]

        hid_map_np_dict_mats = hid_map.log.nparray_dictionary()[exec_id]['matrix']
        hid_map_np_mats = np.array(hid_map.log.nparray()[1][1][5][1:])

        hid_np_dict_vals = xor_hid.log.nparray_dictionary()[exec_id]['value']

        out_map_np_dict_mats = out_map.log.nparray_dictionary()[exec_id]['matrix']
        out_map_np_mats = np.array(out_map.log.nparray()[1][1][5][1:])

        out_np_dict_vals = xor_out.log.nparray_dictionary()[exec_id]['value']

        expected_length = int(len(xor_inputs) * num_epochs / minibatch_size)


        if minibatch_size == 1:
            np.testing.assert_equal(in_np_dict_vals[0:4], xor_inputs[:, None, None, :, :])
        elif minibatch_size == 4:
            np.testing.assert_equal(in_np_vals[0:2][:, :, :], in_np_dict_vals)

        np.testing.assert_equal(in_np_vals, in_np_dict_vals)
        assert in_np_dict_vals.shape == (expected_length, minibatch_size, 1, 1, xor_in.input_shapes)

        assert hid_map_np_dict_mats.shape == (expected_length, xor_in.input_shapes, xor_hid.input_shapes)
        np.testing.assert_equal(hid_map_np_mats, hid_map_np_dict_mats)

        assert hid_np_dict_vals.shape == (expected_length, minibatch_size, 1, 1, xor_hid.input_shapes)

        assert out_map_np_dict_mats.shape == (expected_length, xor_hid.input_shapes, xor_out.input_shapes)
        np.testing.assert_equal(out_map_np_mats, out_map_np_dict_mats)

        assert out_np_dict_vals.shape == (expected_length, minibatch_size, 1, 1, xor_out.input_shapes)

        xor_out.log.print_entries()

    def test_autodiff_loss_tracking(self):
        xor_in = TransferMechanism(name='xor_in',
                                   default_variable=np.zeros(2))

        xor_hid = TransferMechanism(name='xor_hid',
                                    default_variable=np.zeros(10),
                                    function=Logistic())

        xor_out = TransferMechanism(name='xor_out',
                                    default_variable=np.zeros(1),
                                    function=Logistic())

        hid_map = MappingProjection()
        out_map = MappingProjection()

        xor = AutodiffComposition()

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        xor_inputs = np.array(  # the inputs we will provide to the model
            [[0, 0],
             [0, 1],
             [1, 0],
             [1, 1]])

        xor_targets = np.array(  # the outputs we wish to see from the model
            [[0],
             [1],
             [1],
             [0]])

        # train model for a few epochs
        num_epochs = 100
        xor.learn(inputs={"inputs": {xor_in: xor_inputs},
                        "targets": {xor_out: xor_targets},
                        "epochs": num_epochs},
                  execution_mode=pnl.ExecutionMode.PyTorch)

        losses = xor.losses
        # Since the losses track average losses per weight update, and weights are updated every minibatch,
        # and minibatch_size is 1, then there should be num_epochs * num_minibatches = num_epochs * 4
        # total entries
        expected_loss_length = num_epochs * len(xor_inputs)
        assert len(xor.torch_losses) == expected_loss_length

        # test clearing ad losses
        xor.clear_losses(context=xor)
        assert len(xor.losses) == 0

    @pytest.mark.pytorch
    def test_synching_variable_and_value(self):
        """Test synchronization of variables and values of nested Compositions with PyTorch execution."""
        inner_mech_1 = pnl.ProcessingMechanism(name='Inner Mech 1', input_shapes=2, function=pnl.Logistic)
        inner_mech_2 = pnl.ProcessingMechanism(name='Inner Mech 2', input_shapes=3)
        outer_mech_in = pnl.ProcessingMechanism(name='Outer Mech IN', input_shapes=4)
        outer_mech_out = pnl.ProcessingMechanism(name='Outer Mech OUT', input_shapes=5, function=pnl.Logistic)
        inner_comp = pnl.AutodiffComposition(name='Inner Comp',
                                             pathways=[inner_mech_1, inner_mech_2],
                                             # Note: no need to specify synch_node_values_with_torch, since default=RUN
                                             synch_node_variables_with_torch=pnl.RUN)
        outer_comp = pnl.AutodiffComposition(name='Outer Comp',
                                             pathways=[outer_mech_in, inner_comp, outer_mech_out],
                                             # Note: no need to specify synch_node_values_with_torch, since default=RUN
                                             synch_node_variables_with_torch=pnl.RUN)
        targets = outer_comp.infer_backpropagation_learning_pathways(pnl.ExecutionMode.PyTorch)
        result = outer_comp.learn(inputs={outer_mech_in:[[1,2,3,4]],
                                          targets[0]: [[1,1,1,1,1]]},
                                  execution_mode=pnl.ExecutionMode.PyTorch)

        # Outer OUTPUT Mechanism -----------

        # variable
        expected = [[[[5.99972761, 5.99972761, 5.99972761, 5.99972761, 5.99972761]]]]
        # np.testing.assert_allclose(outer_comp.nodes['Outer Mech OUT'].variable, expected)
        np.testing.assert_allclose(outer_comp.nodes['Outer Mech OUT'].parameters.variable.get('Outer Comp'), expected)

        # value
        expected = [[[[0.9975267, 0.9975267, 0.9975267, 0.9975267, 0.9975267]]]]
        # np.testing.assert_allclose(outer_comp.nodes['Outer Mech OUT'].value, expected)
        np.testing.assert_allclose(outer_comp.nodes['Outer Mech OUT'].parameters.value.get('Outer Comp'), expected)

        # Nested INPUT Mechanism  ----------

        # variable
        expected = [[[[10, 10]]]]
        # np.testing.assert_allclose(inner_comp.nodes['Inner Mech 1'].variable, expected)
        # np.testing.assert_allclose(outer_comp.nodes['Inner Comp'].nodes['Inner Mech 1'].variable, expected)
        np.testing.assert_allclose(
            outer_comp.nodes['Inner Comp'].nodes['Inner Mech 1'].parameters.variable.get('Outer Comp'), expected)

        # value
        expected = [[[[0.9999546, 0.9999546]]]]
        # np.testing.assert_allclose(inner_comp.nodes['Inner Mech 1'].value, expected)
        # np.testing.assert_allclose(['Inner Comp'].nodes['Inner Mech 1'].value, expected)
        np.testing.assert_allclose(
            outer_comp.nodes['Inner Comp'].nodes['Inner Mech 1'].parameters.value.get('Outer Comp'), expected)


@pytest.mark.pytorch
class TestBatching:
    def test_call_before_minibatch(self):
        # SET UP MECHANISMS FOR COMPOSITION
        xor_in = TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        # SET UP PROJECTIONS FOR COMPOSITION
        hid_map = MappingProjection(name='hid_map',
                                    matrix=np.random.rand(2, 10),
                                    sender=xor_in,
                                    receiver=xor_hid)

        out_map = MappingProjection(name='out_map',
                                    matrix=np.random.rand(10, 1),
                                    sender=xor_hid,
                                    receiver=xor_out)

        # SET UP COMPOSITION
        xor = AutodiffComposition(learning_rate=10)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # SET UP INPUTS AND TARGETS
        xor_inputs_1 = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
        xor_targets_1 = np.array([[0], [1], [1], [0]])

        # TRAIN COMPOSITION
        inputs_dict_1 = {"inputs": {xor_in: xor_inputs_1},
                         "targets": {xor_out: xor_targets_1},
                         "epochs": 1}

        a = [0]

        def cbm():
            a[0] += 1

        xor.learn(inputs=inputs_dict_1, call_before_minibatch=cbm, execution_mode=pnl.ExecutionMode.PyTorch)

        assert a[0] == 4

    def test_call_after_minibatch(self):
        # SET UP MECHANISMS FOR COMPOSITION
        xor_in = TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        # SET UP PROJECTIONS FOR COMPOSITION
        hid_map = MappingProjection(name='hid_map',
                                    matrix=np.random.rand(2, 10),
                                    sender=xor_in,
                                    receiver=xor_hid)

        out_map = MappingProjection(name='out_map',
                                    matrix=np.random.rand(10, 1),
                                    sender=xor_hid,
                                    receiver=xor_out)

        # SET UP COMPOSITION
        xor = AutodiffComposition(learning_rate=10)

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # SET UP INPUTS AND TARGETS
        xor_inputs_1 = np.array([[0, 0], [0, 1], [1, 0],[1, 1]])
        xor_targets_1 = np.array([[0], [1], [1], [0]])

        # TRAIN COMPOSITION
        inputs_dict_1 = {"inputs": {xor_in: xor_inputs_1},
                         "targets": {xor_out: xor_targets_1},
                         "epochs": 1}

        a = [0]

        def cam():
            a[0] += 1

        xor.learn(inputs=inputs_dict_1, call_after_minibatch=cam, execution_mode=pnl.ExecutionMode.PyTorch)
        assert a[0] == 4

    @pytest.mark.parametrize(
        'eps', (1, 5, 10, 100)
    )
    def test_batching_with_epochs_specified(self, eps):
        # SET UP MECHANISMS FOR COMPOSITION
        xor_in = TransferMechanism(name='xor_in', default_variable=np.zeros(2))
        xor_hid = TransferMechanism(name='xor_hid', default_variable=np.zeros(10), function=Logistic())
        xor_out = TransferMechanism(name='xor_out', default_variable=np.zeros(1), function=Logistic())

        # SET UP PROJECTIONS FOR COMPOSITION
        hid_map = MappingProjection(name='hid_map',
                                    matrix=np.random.rand(2, 10),
                                    sender=xor_in,
                                    receiver=xor_hid)

        out_map = MappingProjection(name='out_map',
                                    matrix=np.random.rand(10, 1),
                                    sender=xor_hid,
                                    receiver=xor_out)

        # SET UP COMPOSITION
        xor = AutodiffComposition(learning_rate=10,
                                  # optimizer_type=opt
                                  )

        xor.add_node(xor_in)
        xor.add_node(xor_hid)
        xor.add_node(xor_out)

        xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
        xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

        # SET UP INPUTS AND TARGETS
        xor_inputs_1 = np.array([[0, 0], [0, 1], [1, 0], [1, 1]])
        xor_targets_1 = np.array([[0], [1], [1], [0]])

        c1 = Context(execution_id='context1')

        # TRAIN COMPOSITION
        inputs_dict_1 = {"inputs": {xor_in: xor_inputs_1},
                         "targets": {xor_out: xor_targets_1},
                         "epochs": eps}

        xor.learn(inputs=inputs_dict_1, context=c1, minibatch_size=2, execution_mode=pnl.ExecutionMode.PyTorch)

        c2 = Context(execution_id='context2')

        xor_inputs_2 = np.array([[0, 0], [0, 1]])
        xor_targets_2 = np.array([[0], [1]])

        inputs_dict_2 = {"inputs": {xor_in: xor_inputs_2},
                         "targets": {xor_out: xor_targets_2},
                         "epochs": 1}

        xor_inputs_3 = np.array([[1, 0], [1, 1]])
        xor_targets_3 = np.array([[1],[0]])

        inputs_dict_3 = {"inputs": {xor_in: xor_inputs_3},
                         "targets": {xor_out: xor_targets_3},
                         "epochs": 1}
        for _ in range(eps):
            xor.learn(inputs=inputs_dict_2, context=c2, minibatch_size=TRAINING_SET, execution_mode=pnl.ExecutionMode.PyTorch)
            xor.learn(inputs=inputs_dict_3, context=c2, minibatch_size=TRAINING_SET, execution_mode=pnl.ExecutionMode.PyTorch)

        c1_results = xor.parameters.results._get(c1)
        c2_results = xor.parameters.results._get(c2)

        np.testing.assert_allclose(c1_results[0][:2], c2_results[0][:2])
        np.testing.assert_allclose(c1_results[0][2:], c2_results[0][2:])

    def test_cross_entropy_loss(self):
        import torch

        m1 = pnl.TransferMechanism()
        p = pnl.MappingProjection()
        m2 = pnl.TransferMechanism()
        adc = pnl.AutodiffComposition(loss_spec=Loss.CROSS_ENTROPY)

        adc.add_linear_processing_pathway([m1, p, m2])
        adc._build_pytorch_representation()

        # The following is based on earlier version of PyTorch that only allows
        # class (index) based target specification (rather than one-hot) for cross-entropy loss:
        # classes = torch.Tensor([2, 1])
        # target = torch.Tensor([1])
        # # Equation for loss taken from https://pytorch.org/docs/stable/nn.html#torch.nn.CrossEntropyLoss
        # assert np.allclose(adc.loss_function(classes, target).detach().numpy(), -1 + np.log(np.exp(2) + np.exp(1)))
        # assert np.allclose(adc.loss_function(output, target).detach().numpy(), -1 + np.log(np.exp(2) + np.exp(1)))

        # Current implementation uses one-hot target specification:
        output = [2,1]
        target = [0,1]

        o = np.array(output)
        t = np.array(target)
        sm = pnl.SoftMax()
        logit = pnl.Logistic()

        # Compute cross-entropy loss as documented for `PyTorchCrossEntropyLoss using target class probabilies vector
        # (see https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html);
        # avoids "inf" value for entries = 0 in target vector
        ce_numpy = -np.sum(np.log( np.exp(o) / np.sum(np.exp(o))) * t)

        # Compute cross-entropy loss using standard mathematical formulation
        #  which can return an "inf" value for entries = 0 in target vector
        ce_pnl = np.sum(pnl.LinearCombination(operation=pnl.CROSS_ENTROPY)([sm(logit(o)),t]))

        output = torch.Tensor(output)
        target = torch.Tensor(target)
        ce_torch = adc.loss_function(output, target).detach().numpy()

        np.testing.assert_allclose(ce_numpy, ce_torch)


@pytest.mark.pytorch
@pytest.mark.parametrize('batch_size', [1, 2, 4])
@pytest.mark.parametrize('batched_results', [False, True])
def test_training_xor_with_batching(batch_size, batched_results):
    """
    This test actually trains an equivalent model in pytorch and compares the losses with the ones from the
    AutodiffComposition learn.
    """

    import torch
    import torch.nn as nn
    import torch.optim as optim

    torch.manual_seed(0)

    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    torch.set_default_dtype(torch.float64)

    HIDDEN_SIZE = 5
    LEARNING_RATE = 0.1
    WEIGHT_DECAY = 0
    NUM_EPOCHS = 2

    X = torch.tensor([[0, 0], [0, 1], [1, 0], [1, 1]], dtype=torch.get_default_dtype()).to(device)
    Y = torch.tensor([[0], [1], [1], [0]], dtype=torch.get_default_dtype()).to(device)

    linear1 = nn.Linear(2, HIDDEN_SIZE, bias=False)
    linear2 = nn.Linear(HIDDEN_SIZE, 1, bias=False)
    sigmoid = nn.Sigmoid()

    model = nn.Sequential(linear1, sigmoid, linear2, sigmoid).to(device) # MLP model

    # Copy the initial weights and biases. Make sure to copy the data, not just the reference.
    weights1 = linear1.weight.detach().cpu().numpy().copy()
    weights2 = linear2.weight.detach().cpu().numpy().copy()
    # bias1 = linear1.bias.detach().numpy().copy() if linear1.bias is not None else None
    # bias2 = linear2.bias.detach().numpy().copy() if linear2.bias is not None else None

    # define cost/loss & optimizer
    # criterion = nn.BCELoss().to(device)
    criterion = nn.MSELoss().to(device)
    optimizer = optim.SGD(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)

    # x, y = X[0:0+batch_size].unsqueeze(1), Y[0:0+batch_size].unsqueeze(1)
    #
    # t = linear1(x); print(f"Linear1 Out: {torch.sum(t).item()}")
    # t = sigmoid(t); print(f"Sigmoid1 Out: {torch.sum(t).item()}")
    # t = linear2(t); print(f"Linear2 Out: {torch.sum(t).item()}")
    # t = sigmoid(t); print(f"Sigmoid2 Out: {torch.sum(t).item()}")
    # loss = criterion(t, y); print(f"Loss: {loss.item()}")

    # A generator to generate batches of X with batch_size, batches will be of shape
    # (batch_size, 1, 2)
    def batcher(x, y):
        for i in range(0, len(x), batch_size):
            yield x[i:i + batch_size].unsqueeze(1), y[i:i + batch_size].unsqueeze(1)

    torch_losses = []
    torch_results = []
    for n in range(NUM_EPOCHS):
        data_loader = batcher(X, Y)

        for step, (x, y) in enumerate(data_loader):
            pred = model(x)
            loss = criterion(pred, y)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            torch_losses.append(loss.item())
            torch_results.append(pred)

    torch_results = torch.stack(torch_results).detach().cpu().numpy()

    autodiff_mode = pnl.ExecutionMode.PyTorch

    xor_in = TransferMechanism(name='xor_in',
                               default_variable=np.zeros(2))

    xor_hid = TransferMechanism(name='xor_hid',
                                default_variable=np.zeros(HIDDEN_SIZE),
                                function=Logistic())

    xor_out = TransferMechanism(name='xor_out',
                                default_variable=np.zeros(1),
                                function=Logistic())

    hid_map = MappingProjection(matrix=weights1.T)
    out_map = MappingProjection(matrix=weights2.T)

    xor = AutodiffComposition(loss_spec=pnl.Loss.MSE, learning_rate=LEARNING_RATE, weight_decay=WEIGHT_DECAY)

    xor.add_node(xor_in)
    xor.add_node(xor_hid)
    xor.add_node(xor_out)

    xor.add_projection(sender=xor_in, projection=hid_map, receiver=xor_hid)
    xor.add_projection(sender=xor_hid, projection=out_map, receiver=xor_out)

    xor_inputs = np.array([[0.0, 0.0],
                           [0.0, 1.0],
                           [1.0, 0.0],
                           [1.0, 1.0]])

    xor_targets = np.array([[0.0], [1.0], [1.0], [0.0]])

    xor.learn(inputs={"inputs": {xor_in: xor_inputs},
                      "targets": {xor_out: xor_targets}},
              epochs=NUM_EPOCHS,
              minibatch_size=batch_size,
              execution_mode=autodiff_mode,
              batched_results=batched_results)

    np.testing.assert_allclose(torch_losses, xor.torch_losses.flatten(), rtol=1e-5)

    if batched_results:
        np.testing.assert_allclose(torch_results, xor.results, rtol=1e-5)
    else:
        torch_results_unbatched = torch_results.reshape( (NUM_EPOCHS * (len(X) // batch_size) * batch_size, 1, 1) )
        np.testing.assert_allclose(torch_results_unbatched, xor.results, rtol=1e-5)
