Asserts that key is valid NestedMap key.

Reads from a Hyperparams proto.

Generates batch_size sets of points, the k-th example in the batch has n-k centers (i-th center is on (i, i)) and m points near each center.

Return the number of leading combined-layers dims for CreateVariable().

Returns new state with leading shape=[batch, max_steps].

Maximum rotation angle in world coordinates.

Parameters for this task scheduler.

Quantizes a tensor with a known, natural range. Args: t: Tensor to quantize. min_value: Min value of the range. max_value: Max

Resets step_seed to specified value.

Load the checkpoint from specified path.

Set the logdir member in all MetricHistory.params objects in params. Args: params: global model params. logdir: root dir for current

Sets the current value of the accumulator Tensor.

Streams t steps. Args: theta: A NestedMap of layer params. inputs: [b, t, d]. paddings: A 0/1 valued tensor of shape [b, t].

Computes cross entropy.

If the given func_name is a Defun, add its sub-nodes into nodes.

Bias layer. The bias is added to the last dimension of the input.

Input generator params for Librispeech.

Create different program param instance per inputs.

Exports metrics externally.

y = fn(x). Applies a fn: tuple(Tensor) -> a single Tensor or tuple(Tensor) to the input tuple. Typically, fn is a very simple python functio

Graph Isomorphism Network (GIN) [1]. This implements the GIN network. The input goes through multiple GINIntermediateLayers which performs th

Returns a `.NestedMap` with x[index, :] for each tensor x in nmap. Args: nmap: A `.NestedMap` of tensors. index: A tf scalar integer. Perfo

Construct a job spec param with the given number of replicas.

Create empty TensorArray which can store max_seq_length elements.

Remove a variable by name from `layer`. This is usually used to avoid copying the variable to TPU, for example, by the tf.cast when accessing lay

Updates t-th row in accumulators. Args: nmap_acc: A `.NestedMap` of tensors. The accumulators. nmap_x: A `.NestedMap` of tensors. The updat

Returns current context or None.

Applies clipping to x. Args: theta: Layer theta. x: Input tensor to clip. **kwargs: Additional implementation specific kwargs.

Copies BaseLayer params from `from_params` to `to_params`.

Create a matrix of coordinate locations corresponding to a dense grid. Example: To create (x, y) coordinates corresponding over a 10x10 grid with

Transformer decoder layer, extend one step in autoregressive decoding. query_vec and aux_* may have different batch sizes, e.g., during a beam

Apply the stacking to inputs along the time axis. Args: inputs: The inputs tensor. It is expected to be of shape [batch, time, feat

Compute features for the pillars and convert them back to a dense grid. Args: theta: A `.NestedMap` object containing variable values of th

Converts every dict in nested structure 'x' to a NestedMap.

(inputs, paddings) -> (encoded, paddings). Args: name: the string name of the encoder layer params. stride: To pool every k tokens, s

Recursively gets values of all accumulators. Returns: `.NestedMap` of Tensors for each registered accumulator.

Returns a copy of self containing only entries matching `keys`. Note: This function permits lingvo-nestedmap-style nested keys such as `m.f

Takes int32 token ids and returns approximate detokenized strings.

Packs 'values' according to 'tmpl'.

Post-processes decoder out and updates contents of `decode_metrics_dict`. Args: decode_out_dict: A dictionary of Tensors fetched. Could als

Called once for each train loop. BaseModel.ProcessFPropResults is also called on each loop, so you can put your implementation wherever it is

Restore from latest checkpoint if available, or initialize.

Performs target sequence sampling. Args: theta: A NestedMap object containing weights' values of this layer and its children layers

Save the checkpoint. Args: sess: tf.Session. gsteps: Current global step. sync: Whether to wait for the saving operations to fi

Scales gradients according to training params. Args: var_grads: a `.NestedMap` whose values are (var, grad) pairs. gradient_adjuster:

Returns whether the trial should stop.

Computes the value vector given the query of the current step. This differs from ExtendStep() which requires key/value seq lengths being know

Converts 'shape' to a static shape.