Set the initial step.

Setup the assets like dataset, pretrained models, etc.

Setup the networks in the model for distributed computation. It automactically gather all modules.

Shuffle `arr` along `axis`. Args: arr: (*,) axis: Returns: (*,)

Returns a dictionary that can be saved or load.

Returns a dictionary that can be saved or load.

Returns a dictionary that can be saved or load.

Returns a dictionary that can be saved or load.

Given x y z index, change to the linear index. (matlab's sub2ind) Args: idx: (*, n, 3) int64_t size: (*, 3) int64_t Returns: (*, n) int64_t

This function exposes the underlying subspace operators. See :py:class:`Subspace` for detailed documentation.

Args: xls (*, style_dim) or list of (*, style_dim) or None: input tensor on the right hand side xrs (

Setup before test loop.

Setup before training loop.

Sample the feature map at the uv values. Args: uv: (b, *p, 2) values between [0, 1] feature_map: (b

Setup before validation loop.

query the texture map at locations uv Args: uv: (*, 2) u is in the x/width direction, v is in the y/height direction,

Returns: input_rgbd_images: RGBDImage, (b, q, h, w), ray: Ray, (b, q=n_target_img, ho

Return the ith's sample (in global index) as if calling samples[idx]. Note: 1. Each sample is returned in their native shape (no

Returns the content of the i-th line in the index file.

Returns: point_cloud: PointCloud, (b=1, n), n: number of points ray: Ray, (b=1, q=n_

slice the camera in the b dimension. Always retain (b, q, 4, 4) even when ib is int.

Create the logger. Args: log_filename: Filename to save the log. None: do not save to a file. Re

Args: convert_to_float: whether to convert input values (from Tensor, ndarray) to float

Args: exp_tag (str): name of the experiment config_filename (str): the yaml or json f

r""" Vary the learning rate every training step according to: .. math:: lr = \text{factor} * d_{model}^{-0.5} * \

Create memmaps that store chunked of the samples. Args: working_dir: the directory where all the chunked memmaps

Args: index_filename: the index file that lists all the data, one per line. with_uid:

Args: batch_size: typical batch_size collate_fn: collate function used to combine sam

Args: batch_size: batch size drop_last: whether to drop the last few data that cannot

Args: seq_lens: sequence length of individual samples. batch_size: batch size of indi

Args: batch_size: batch size max_batch_combined_size: limitation on batch_size * seq_

Create an Alex Graves' model. Args: param_dict: A :py:class:`ParamGraves` object to define the hyper-parameters o

Create a Variational Recurrent Neural Network (VRNN) model. Args: param_dict: A :py:class:`ParamVRNN` ob

Convenient helper nn.Module to create a stack of linear layers. Args: num_layers: Total number of linear

r""" Args: in_features (int): input feature dimension out_features (int): output featu

Construct multiple stacked conv1D layers. It applies a low-pass filtering for every upsampling/downsampling to avoid aliasing.

r""" Args: in_channels (int): number of input channels of the blur_kernel out_channels (int):

Given input, the layer computes a mixture of 1D gaussians as the attention weights. Args: num_mixtures:

See :py:class:`TacotronContentEncoder` for details.

r""" Args: in_features (int): input feature dimension out_features (int): output featu

r"""Construct multiple stacked modulated subspace layers sandwiched between nonlinearity. Args: in_features (int):

Construct nested dropout layer, which drops the last dimension. Note that it creates a mask of shape (B, C), so if the input x

let :math:`p_t = p` if target = 1, :math:`1-p` if target != 1 let :math:`a_t = a` if target = 1, :math:`1-p` if target != 1

Create a one-dimensional transposed convolution layer whose weights are initialized with a chosen method. The usage of the l

Multiple Conv1D layers with normalization and nonlinearity. Args: num_layers (int): total number of laye

A multi-layer RNN composed of LSTMCells (or the provided lstm_cell_fn). Args: num_layers: Number of RNN

Args: texture: (h, w, dim) for example, an rgb image, a displacement map, a bump map, etc mode:

Args: u: (*, 3) the mean direction. (*) is the dimension of independent components`.

Args: mesh: class Mesh containing cleaned and scaled mesh and ray_casting scene. n_target_imgs:

Args: scene_dir: the root dir of the scene, e.g., /mnt/task_runtime/trove/ml-hypersim-1.0.0/data

Args: mode: n_imgs: number of cameras in a set total: total number of se

Apply the color correction to an rgbd_image Args: correction_type: 'wrgb': the correction is 3 scalars \

- randomly select sample indexes to be used in the epoch - reset the batch_sampler :return: