Support layer custom names

trax/layers/base.py

@@ -85,15 +85,17 @@ class Layer(object):
   outputs are spliced back into the stack.
   """
 
-  def __init__(self, n_in=1, n_out=1):
+  def __init__(self, n_in=1, n_out=1, name=None):
     """Creates a partially initialized, unconnected layer instance.
 
     Args:
       n_in: Number of inputs expected by this layer.
       n_out: Number of outputs promised by this layer.
+      name: Descriptive name for this layer.
     """
     self._n_in = n_in
     self._n_out = n_out
+    self._name = name or self.__class__.__name__
     self._sublayers = ()  # Default is no sublayers.
     self._input_signature = None
     self._rng = None
@@ -109,7 +111,7 @@ def __init__(self, n_in=1, n_out=1):
     self._jit_cache = {}
 
   def __repr__(self):
-    class_str = self.__class__.__name__
+    class_str = self._name
     fields_str = 'in={},out={}'.format(self.n_in, self.n_out)
     objs = self.sublayers
     if objs:
@@ -316,7 +318,7 @@ def init(self, input_signature, rng=None):
       else:
         return (EMPTY_WEIGHTS, state)
     except Exception as e:
-      name, trace = self.__class__.__name__, _short_traceback(skip=3)
+      name, trace = self._name, _short_traceback(skip=3)
       raise LayerError(name, 'init', self._caller,
                        input_signature, trace) from e
 
@@ -446,7 +448,7 @@ def pure_fn(self, x, weights, state, rng):
       return outputs, s
 
     except Exception as e:
-      name, trace = self.__class__.__name__, _short_traceback()
+      name, trace = self._name, _short_traceback()
       raise LayerError(name, 'pure_fn',
                        self._caller, signature(x), trace) from e
 
@@ -480,7 +482,7 @@ def call_on_input(x, weights, state, rng):
           input_signature, weight_signature, self.state, rng)
       return s
     except Exception as e:
-      name, trace = self.__class__.__name__, _short_traceback(skip=3)
+      name, trace = self._name, _short_traceback(skip=3)
       raise LayerError(name, '_forward_abstract', self._caller, input_signature,
                        trace) from e
 
@@ -564,15 +566,15 @@ def vjpfun(grad):
     return output, state
 
 
-def layer(n_in=1, n_out=1, new_weights_fn=None):
+def layer(n_in=1, n_out=1, new_weights_fn=None, name=None):
   """Returns a decorator that converts a function into a Layer class builder."""
 
   def _build_layer_class(raw_fn):
     """Returns a Layer class whose callable instances execute the function."""
 
     def _init(self, **kwargs):
       self._kwargs = kwargs  # pylint: disable=protected-access
-      Layer.__init__(self, n_in=n_in, n_out=n_out)
+      Layer.__init__(self, n_in=n_in, n_out=n_out, name=name)
 
     def _forward(self, x, weights):
       """Uses this layer as part of a forward pass through the model."""

trax/layers/base_test.py

@@ -195,6 +195,28 @@ def forward(self, inputs, weights, **kwargs):
     layer(0, n_accelerators=1)
     layer(0, n_accelerators=1)
 
+  def test_custom_name(self):
+    layer = base.Layer()
+    self.assertIn('Layer', str(layer))
+    self.assertNotIn('CustomLayer', str(layer))
+
+    layer = base.Layer(name='CustomLayer')
+    self.assertIn('CustomLayer', str(layer))
+
+    @base.layer()
+    def DefaultDecoratorLayer(x, **unused_kwargs):
+      return x
+
+    layer = DefaultDecoratorLayer()  # pylint: disable=no-value-for-parameter
+    self.assertIn('DefaultDecoratorLayer', str(layer))
+
+    @base.layer(name='CustomDecoratorLayer')
+    def NotDefaultDecoratorLayer(x, **unused_kwargs):
+      return x
+
+    layer = NotDefaultDecoratorLayer()  # pylint: disable=no-value-for-parameter
+    self.assertIn('CustomDecoratorLayer', str(layer))
+
 
 if __name__ == '__main__':
   absltest.main()

trax/layers/combinators.py

@@ -41,8 +41,8 @@ class Serial(base.Layer):
   1-input 1-output no-op.
   """
 
-  def __init__(self, *sublayers):
-    super(Serial, self).__init__()
+  def __init__(self, *sublayers, name=None):
+    super(Serial, self).__init__(name=name)
 
     sublayers = _ensure_flat(sublayers)
     self._sublayers = sublayers
@@ -191,17 +191,18 @@ class Parallel(base.Layer):
   following input(s).
   """
 
-  def __init__(self, *sublayers):
+  def __init__(self, *sublayers, name=None):
     """The constructor.
 
     Args:
       *sublayers: A list of sublayers.
+      name: Descriptive name for this layer.
 
     Returns:
       A new layer in which each of the given sublayers applies to its
       corresponding span of elements in the dataflow stack.
     """
-    super(Parallel, self).__init__()
+    super(Parallel, self).__init__(name=name)
     sublayers = self._validate(sublayers)
     self._n_layers = len(sublayers)
     self._sublayers = sublayers
@@ -437,7 +438,7 @@ def new_weights_and_state(self, input_signature):
     return self.sublayer.init(layer_signature)
 
 
-def Branch(*layers):
+def Branch(*layers, name='Branch'):
   """Combinator that applies a list of layers in parallel to copies of inputs.
 
   Each layer in the input list is applied to as many inputs from the stack
@@ -459,6 +460,7 @@ def Branch(*layers):
 
   Args:
     *layers: list of layers
+    name: Descriptive name for this layer.
 
   Returns:
     the branch layer
@@ -467,7 +469,7 @@ def Branch(*layers):
     return layers[0]
   parallel_layer = Parallel(*layers)
   indices = [list(range(layer.n_in)) for layer in parallel_layer.sublayers]
-  return Serial(Select(_deep_flatten(indices)), parallel_layer)
+  return Serial(Select(_deep_flatten(indices)), parallel_layer, name=name)
 
 
 def Residual(*layers, **kwargs):
@@ -513,7 +515,7 @@ def Swap(xs, **unused_kwargs):
   return (xs[1], xs[0])
 
 
-def Select(indices, n_in=None):
+def Select(indices, n_in=None, name=None):
   """Copies, reorders, or deletes stack elements according to `indices`.
 
   Args:
@@ -522,6 +524,7 @@ def Select(indices, n_in=None):
     n_in: Number of input elements to pop from the stack, and replace with
         those specified by `indices`. If not specified, its value will be
         calculated as `max(indices) + 1`.
+    name: Descriptive name for this layer.
 
   Returns:
     Tensors, matching the number selected (`n_out = len(indices)`).
@@ -533,7 +536,7 @@ def Select(indices, n_in=None):
   if n_in is None:
     n_in = max(indices) + 1
 
-  @base.layer(n_in=n_in, n_out=len(indices))
+  @base.layer(n_in=n_in, n_out=len(indices), name=name)
   def Selection(xs, **unused_kwargs):  # pylint: disable=invalid-name
     if not isinstance(xs, (tuple, list)):
       xs = (xs,)

trax/layers/combinators_test.py

@@ -78,6 +78,13 @@ def some_layer():
     output_shape = base.check_shape_agreement(layer, input_signature)
     self.assertEqual(output_shape, expected_shape)
 
+  def test_serial_custom_name(self):
+    layer = cb.Serial(cb.Dup(), cb.Dup())  # pylint: disable=no-value-for-parameter
+    self.assertIn('Serial', str(layer))
+
+    layer = cb.Serial(cb.Dup(), cb.Dup(), name='Branch')  # pylint: disable=no-value-for-parameter
+    self.assertIn('Branch', str(layer))
+
   def test_branch_noop_dup(self):
     layer = cb.Branch([], cb.Dup())
     input_signature = ShapeDtype((3, 2))
@@ -99,6 +106,10 @@ def test_branch_one_layer(self):
     output_shape = base.check_shape_agreement(layer, input_signature)
     self.assertEqual(output_shape, expected_shape)
 
+  def test_branch_name(self):
+    layer = cb.Branch(cb.Add(), divide_by(0.5))  # pylint: disable=no-value-for-parameter
+    self.assertIn('Branch', str(layer))
+
   def test_select_computes_n_in(self):
     layer = cb.Select([0, 0])
     self.assertEqual(layer.n_in, 1)
@@ -150,6 +161,13 @@ def test_parallel_no_ops(self):
     output_shape = base.check_shape_agreement(layer, input_signature)
     self.assertEqual(output_shape, expected_shape)
 
+  def test_parallel_custom_name(self):
+    layer = cb.Parallel(cb.Dup(), cb.Dup())  # pylint: disable=no-value-for-parameter
+    self.assertIn('Parallel', str(layer))
+
+    layer = cb.Parallel(cb.Dup(), cb.Dup(), name='DupDup')  # pylint: disable=no-value-for-parameter
+    self.assertIn('DupDup', str(layer))
+
   def test_drop(self):
     layer = cb.Drop()
     input_signature = ShapeDtype((3, 2))