Preserve order in ChainedOutcomeTransform (#440)

Summary:
Pull Request resolved: #440

Using OrderedDict, perserves the order. Otherwise, ModuleDict orders the transforms alphanumerically by key.

Reviewed By: Balandat

Differential Revision: D21477968

fbshipit-source-id: aa85d450ccb9b8baa651f2a82df526e91e4daa69

Author: sdaulton

botorch/models/transforms/outcome.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections import OrderedDict
 from typing import List, Optional, Tuple
 
 import torch
@@ -87,7 +88,7 @@ def __init__(self, **transforms: OutcomeTransform) -> None:
                 kwargs are used as the keys for accessing the individual
                 transforms on the module.
         """
-        super().__init__(transforms)
+        super().__init__(OrderedDict(transforms))
 
     def forward(
         self, Y: Tensor, Yvar: Optional[Tensor] = None

test/models/transforms/test_outcome.py

@@ -328,11 +328,11 @@ def test_chained_outcome_transform(self):
             # test init
             tf1 = Log()
             tf2 = Standardize(m=m, batch_shape=batch_shape)
-            tf = ChainedOutcomeTransform(log=tf1, standardize=tf2)
+            tf = ChainedOutcomeTransform(b=tf1, a=tf2)
             self.assertTrue(tf.training)
-            self.assertEqual(sorted(tf.keys()), ["log", "standardize"])
-            self.assertEqual(tf["log"], tf1)
-            self.assertEqual(tf["standardize"], tf2)
+            self.assertEqual(list(tf.keys()), ["b", "a"])
+            self.assertEqual(tf["b"], tf1)
+            self.assertEqual(tf["a"], tf2)
 
             # make copies for validation below
             tf1_, tf2_ = deepcopy(tf1), deepcopy(tf2)