def Params(cls):
"""Returns the layer params."""
p = hyperparams.Params()
p.Define('cls', cls, 'Cls that this param object is associated with.')
p.Define('inference_driver_name', cls._INFERENCE_DRIVER_NAME,
'Name of the inference driver used to construct this layer.')
p.Define('name', '', 'Name of this layer object.')
p.Define('dtype', tf.float32, 'Datatype to use.')
# None value will make FProp use dtype instead of fprop_dtype.
# TODO(lepikhin): all @function.Defun should use p.fprop_dtype if it is set.
p.Define('fprop_dtype', None, 'Activations datatype to use.')
p.Define(
'random_seed', None,
'Random seed for deterministic unittests. This '
'is inherited by child layers if they do not set a random_seed.')
p.Define('vn', DefaultVN(), 'How variational noise should be applied.')
p.Define('params_init', py_utils.DefaultParamInit(),
'How params should be initialized.')
# is_eval is used to generate graph for eval purpose, typically
# the eval graph is forward pass of training graph without
# regularization, e.g. dropout.
p.Define('is_eval', None, 'True if in eval mode.')
# In addition to is_eval, also makes additional alterations for graphs
# being used for inference.
p.Define('is_inference', None, 'True if in inference mode.')
# In addition to is_eval/is_inference, indicate that the inference graph is
# for a single step.
p.Define(
'allow_implicit_capture', None,
'When using Defuns, code often asserts that the Defun does not '
'capture undeclared inputs. This eliminates a source of bugs '
'at the expense of making some kinds of models or utilities '
'hard/impossible to use. Setting this to True/False (versus None) '
'causes the setting to apply to this layer and its children.')
return p
def Params(cls: Type[BaseLayerT]) -> BaseLayerParamsT:
"""Returns the layer params."""
p = hyperparams.InstantiableParams(cls)
p.Define('inference_driver_name', cls._INFERENCE_DRIVER_NAME,
'Name of the inference driver used to construct this layer.')
p.Define('name', '', 'Name of this layer object.')
p.Define('dtype', tf.float32, 'Datatype to use.')
# None value will make FProp use dtype instead of fprop_dtype.
# TODO(lepikhin): all @tf.Defun should use p.fprop_dtype if it is set.
p.Define('fprop_dtype', None, 'Activations datatype to use.')
p.Define(
'random_seed', None,
'Random seed for deterministic unittests. This '
'is inherited by child layers if they do not set a random_seed.')
p.Define('vn', py_utils.DefaultVN(),
'How variational noise should be applied.')
p.Define(
'params_init', py_utils.DefaultParamInit(),
'How model weights should be initialized. Not to be confused with '
'hyperparams.')
p.Define('add_name_to_theta', False,
'Wrap theta with tf.identity(var_name).')
# Makes additional alterations for graphs being used for inference.
p.Define('is_inference', None, 'True if in inference mode.')
# In addition to is_inference, indicate that the inference graph is
# for a single step.
p.Define(
'allow_implicit_capture', None,
'When using Defuns, code often asserts that the Defun does not '
'capture undeclared inputs. This eliminates a source of bugs '
'at the expense of making some kinds of models or utilities '
'hard/impossible to use. Setting this to True/False (versus None) '
'causes the setting to apply to this layer and its children.')
p.Define(
'skip_lp_regularization', None,
'If True, all variables in this layer will skip Lp regularization. '
'If None/False, only variables explicitly in the '
'SKIP_LP_REGULARIZATION collection will skip Lp regularization. '
'Also propagated to child layers with default settings (None).')
# SPMD partition related params.
p.Define(
'device_mesh', None,
'A numpy.ndarray specifying the topology of a device mesh to place the'
' computations onto. If device_mesh is None, it is assumed to be a'
' single device. Here are some examples:'
' np.array([0, 1, 2, 3, 4, 5, 6, 7]) which is a 1d mesh with 8 devices,'
' np.array([[0, 1, 2, 3], [4, 5, 6, 7]]) which is 2d matrix of 8'
' devices.')
p.Define(
'weight_split_dims_mapping', None,
'Relevant only if device_mesh above is not None. If not None, it '
'specifies how weight of this layer or those of the sublayers should '
'be sharded over device mesh. ')
p.Define(
'activation_split_dims_mapping', None,
'Relevant only if device_mesh above is not None. If not None, it '
'specifies how activation of this layer or those of the sublayers '
'should be sharded over device mesh. ')
return p
def Params(cls):
"""Returns the layer params."""
p = hyperparams.InstantiableParams(cls)
p.Define('inference_driver_name', cls._INFERENCE_DRIVER_NAME,
'Name of the inference driver used to construct this layer.')
p.Define('name', '', 'Name of this layer object.')
p.Define('dtype', tf.float32, 'Datatype to use.')
# None value will make FProp use dtype instead of fprop_dtype.
# TODO(lepikhin): all @tf.Defun should use p.fprop_dtype if it is set.
p.Define('fprop_dtype', None, 'Activations datatype to use.')
p.Define(
'random_seed', None,
'Random seed for deterministic unittests. This '
'is inherited by child layers if they do not set a random_seed.')
p.Define('vn', py_utils.DefaultVN(),
'How variational noise should be applied.')
p.Define(
'params_init', py_utils.DefaultParamInit(),
'How model weights should be initialized. Not to be confused with '
'hyperparams.')
# Makes additional alterations for graphs being used for inference.
p.Define('is_inference', None, 'True if in inference mode.')
# In addition to is_inference, indicate that the inference graph is
# for a single step.
p.Define(
'allow_implicit_capture', None,
'When using Defuns, code often asserts that the Defun does not '
'capture undeclared inputs. This eliminates a source of bugs '
'at the expense of making some kinds of models or utilities '
'hard/impossible to use. Setting this to True/False (versus None) '
'causes the setting to apply to this layer and its children.')
p.Define(
'skip_lp_regularization', None,
'If True, all variables in this layer will skip Lp regularization. '
'If None/False, only variables explicitly in the '
'SKIP_LP_REGULARIZATION collection will skip Lp regularization. '
'Also propagated to child layers with default settings (None).')
return p
def testIsDefaultParamInit(self): p = py_utils.DefaultParamInit() self.assertTrue(py_utils.IsDefaultParamInit(p))
