Defines a parameter. Args: name: The parameter name. Must only contain lowercase letters, numbers, and underscores. Must start with

Test session context manager.

Evaluates `function_or_fetches` with `feed_dict` as input. Args: function_or_fetches: Eager tensors, or functions decorated by @Def

Calls `FProp` with this layer's parameters.

Create a sub layer. The created sub layer can be accessed by `name`. E.g.:: self.CreateChild('foo', foo_params) self.foo.FProp..

Creates a deep copy of self.

Create a variable of this layer according to the parameter `var_params`. E.g.:: def __init__(self, ...): # A layer's constructor

Returns the size (num of elements) of the tensor.

Forward function. step_inputs, state0, step_outputs, and state1 should each be a `.NestedMap` of tensor values. Each tensor must be of shape

scale * tf.random.normal(0, 1.0).

Returns a list containing the flattened values in the `.NestedMap`. Unlike py_utils.Flatten(), this will only descend into lists, dicts, and

Returns a copy of this `.NestedMap` with fn applied on each value.

scale.

scale * tf.random.uniform(-1.0, 1.0).

Applies stride to the inputs. Args: theta: weights defined in this layer. x: input tensor, [..., time, ...]. Stride is applied to the

Xavier initialization (x = sqrt(6. / (in + out)); [-x, x]).

Look up style embedding.

Gathers values from v. Args: v: A tensor of shape [B, T, N, D] indx: A tensor of shape [B, N, K, W] Returns: A v

Params for the Scones output layer.

Encodes source as represented by `inputs` and `paddings`. Args: theta: A `.NestedMap` object containing weights' values of this layer and

Looks up embedding vectors for ids. Args: theta: Named tuple with the weight matrix for the embedding. ids: A rank-N int32 tensor.

Connects sub-layers into a data flow graph.

Return the task associated with 'task_name'. Args: task_name: string, the name of the model task to be returned. Returns: An ins

Create a list or dict of sub layers. The created sub layer list can be accessed by `name`. E.g.:: self.CreateChildren('foo', ...)

Flatten the `.NestedMap` and returns <key, value> pairs in a list. Returns: A list of <key, value> pairs, where keys for nested entries wil

Returns a copy of this with each value replaced by a value in lst.

Deep-copies the structure but not the leaf objects.

Connects sub-layers sequentially.

Constructs a `Params` object for given model and dataset, obeying flags. In case of default model, params may be updated based on the flags `

Returns current context or None.

A wrapper of InitForSourcePacked of child attention layer. Args: theta: A `.NestedMap` object containing weights' values of this la

Configs for TransformerStack.

Transform the points in bounding box `i`.

Updates the metrics. Args: label: An array to specify the groundtruth binary labels. Values must be either 0 or 1. prob: An a

Reshapes `tensors`, collapsing the leading `batch_shape` to one dim. The shapes of all `tensors` must begin with `batch_shape`. Args:

Computes the context vector given the current query output. Unlike `ComputeContextVectorWithSource` which explicitly asks for the packed sour

Quantizes an activation. act_name must have been previously created via TrackQActs. Args: act_name: Previously created act_name to qua

Sets sharding attribute for the Tensor. Split across dim=0.

Returns tensor's shape as a list which can be unpacked, unlike tf.shape. If the tensor is unranked, and ndims is None, returns the shape as a Tenso

Initialize this datasource. Args: input_generator: A reference to the external input_generator.

Quantizes a weight. Args: w: The weight tensor. domain: Name of the QDomain to use for quantization. Returns: The weights q

Defaults params for input generators.

Returns a list of human-interpretable strings.

Generates a JTensor of sinusoids with different frequencies for a step. Args: theta: A `.NestedMap` object containing weights' values of th

Computes xent loss given the language model input activations. Args: theta: A `.NestedMap` object containing weights' values of this layer

Pythonic dict-like iteration.

Expects a NestedMap as input and produces the value for the given key.

Parameters of this class.

Returns the current epoch number. Epoch number starts from 1 and reflects the epoch number of the record returned by the next Yield() call.