SYMBOL INDEX (339 symbols across 14 files)

FILE: complex/complex_functions.py
  function complex_matmul (line 10) | def complex_matmul(A, B):
  function complex_avg_pool1d (line 20) | def complex_avg_pool1d(input, *args, **kwargs):
  function complex_avg_pool2d (line 26) | def complex_avg_pool2d(input, *args, **kwargs):
  function complex_normalize (line 35) | def complex_normalize(input):
  function complex_leaky_relu (line 45) | def complex_leaky_relu(input, negative_slope):
  function complex_relu (line 48) | def complex_relu(input):
  function complex_sigmoid (line 51) | def complex_sigmoid(input):
  function complex_tanh (line 54) | def complex_tanh(input):
  function complex_opposite (line 57) | def complex_opposite(input):
  function complex_stack (line 60) | def complex_stack(input, dim):
  function _retrieve_elements_from_indices (line 65) | def _retrieve_elements_from_indices(tensor, indices):
  function _retrieve_elements_from_indices3d (line 70) | def _retrieve_elements_from_indices3d(tensor, indices):
  function complex_upsample (line 75) | def complex_upsample(input, size=None, scale_factor=None, mode='nearest',
  function complex_upsample2 (line 87) | def complex_upsample2(input, size=None, scale_factor=None, mode='nearest',
  function complex_max_pool2d (line 102) | def complex_max_pool2d(input,kernel_size, stride=None, padding=0,
  function complex_max_pool3d (line 126) | def complex_max_pool3d(input,kernel_size, stride=None, padding=0,
  function complex_adaptive_avg_pool3d (line 151) | def complex_adaptive_avg_pool3d(input,kernel_size, stride=None, padding=0,
  function complex_dropout (line 176) | def complex_dropout(input, p=0.5, training=True):
  function complex_dropout2d (line 186) | def complex_dropout2d(input, p=0.5, training=True):
  function complex_dropout3d (line 195) | def complex_dropout3d(input, p=0.5, training=True):

FILE: complex/complex_layers.py
  function apply_complex (line 20) | def apply_complex(fr, fi, input, dtype = torch.complex64):
  class ComplexDropout (line 24) | class ComplexDropout(Module):
    method __init__ (line 25) | def __init__(self,p=0.5):
    method forward (line 29) | def forward(self,input):
  class ComplexDropout2d (line 35) | class ComplexDropout2d(Module):
    method __init__ (line 36) | def __init__(self,p=0.5):
    method forward (line 40) | def forward(self,input):
  class ComplexDropout3d (line 46) | class ComplexDropout3d(Module):
    method __init__ (line 47) | def __init__(self,p=0.5):
    method forward (line 51) | def forward(self,input):
  class ComplexMaxPool2d (line 57) | class ComplexMaxPool2d(Module):
    method __init__ (line 59) | def __init__(self,kernel_size, stride= None, padding = 0,
    method forward (line 69) | def forward(self,input):
  class ComplexMaxPool3d (line 75) | class ComplexMaxPool3d(Module):
    method __init__ (line 77) | def __init__(self,kernel_size, stride=None, padding = 0,
    method forward (line 87) | def forward(self,input):
  class ComplexAvgPool2d (line 94) | class ComplexAvgPool2d(Module):
    method __init__ (line 96) | def __init__(self,kernel_size, stride= None, padding = 0,
    method forward (line 106) | def forward(self,input):
  class ComplexReLU (line 113) | class ComplexReLU(Module):
    method forward (line 115) | def forward(self,input):
  class ComplexSigmoid (line 118) | class ComplexSigmoid(Module):
    method forward (line 120) | def forward(self,input):
  class ComplexTanh (line 123) | class ComplexTanh(Module):
    method forward (line 125) | def forward(self,input):
  class ComplexConvTranspose2d (line 128) | class ComplexConvTranspose2d(Module):
    method __init__ (line 130) | def __init__(self,in_channels, out_channels, kernel_size, stride=1, pa...
    method forward (line 141) | def forward(self,input):
  class ComplexConv2d (line 144) | class ComplexConv2d(Module):
    method __init__ (line 146) | def __init__(self,in_channels, out_channels, kernel_size=3, stride=1, ...
    method forward (line 152) | def forward(self,input):
  class ComplexConv3d (line 155) | class ComplexConv3d(Module):
    method __init__ (line 157) | def __init__(self, in_channels, out_channels, kernel_size=3, stride=1,...
    method forward (line 163) | def forward(self,input):
  class ComplexLinear (line 167) | class ComplexLinear(Module):
    method __init__ (line 169) | def __init__(self, in_features, out_features):
    method forward (line 174) | def forward(self, input):
  class NaiveComplexBatchNorm1d (line 178) | class NaiveComplexBatchNorm1d(Module):
    method __init__ (line 182) | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, \
    method forward (line 188) | def forward(self,input):
  class NaiveComplexBatchNorm2d (line 191) | class NaiveComplexBatchNorm2d(Module):
    method __init__ (line 195) | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, \
    method forward (line 201) | def forward(self,input):
  class NaiveComplexBatchNorm3d (line 204) | class NaiveComplexBatchNorm3d(Module):
    method __init__ (line 208) | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, \
    method forward (line 214) | def forward(self,input):
  class NaiveComplexLayerNorm (line 217) | class NaiveComplexLayerNorm(Module):
    method __init__ (line 221) | def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
    method forward (line 226) | def forward(self,input):
  class _ComplexBatchNorm (line 229) | class _ComplexBatchNorm(Module):
    method __init__ (line 231) | def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True,
    method reset_running_stats (line 257) | def reset_running_stats(self):
    method reset_parameters (line 265) | def reset_parameters(self):
  class ComplexBatchNorm2d (line 272) | class ComplexBatchNorm2d(_ComplexBatchNorm):
    method forward (line 274) | def forward(self, input):
  class ComplexBatchNorm1d (line 345) | class ComplexBatchNorm1d(_ComplexBatchNorm):
    method forward (line 347) | def forward(self, input):
  class ComplexGRUCell (line 419) | class ComplexGRUCell(Module):
    method __init__ (line 424) | def __init__(self, input_length=10, hidden_length=20):
    method reset_gate (line 443) | def reset_gate(self, x, h):
    method update_gate (line 450) | def update_gate(self, x, h):
    method update_component (line 456) | def update_component(self, x, h, r):
    method forward (line 462) | def forward(self, x, h):
  class ComplexBNGRUCell (line 477) | class ComplexBNGRUCell(Module):
    method __init__ (line 482) | def __init__(self, input_length=10, hidden_length=20):
    method reset_gate (line 503) | def reset_gate(self, x, h):
    method update_gate (line 510) | def update_gate(self, x, h):
    method update_component (line 516) | def update_component(self, x, h, r):
    method forward (line 522) | def forward(self, x, h):

FILE: complex/complex_module.py
  function apply_complex (line 9) | def apply_complex(F_r, F_i, X):
  function apply_complex_sep (line 13) | def apply_complex_sep(F_r, F_i, X):
  function complex_mul (line 18) | def complex_mul(X, Y):
  function complex_bmm (line 26) | def complex_bmm(X, Y):
  function complex_softmax (line 34) | def complex_softmax(X):
  function transpose_qkv (line 39) | def transpose_qkv(x, num_heads: int):
  function transpose_output (line 45) | def transpose_output(x, num_heads: int):
  class ComplexDropout (line 51) | class ComplexDropout(nn.Module):
    method __init__ (line 52) | def __init__(self, p=0.5):
    method forward (line 56) | def forward(self, X):
  class ComplexGELU (line 64) | class ComplexGELU(nn.Module):
    method __init__ (line 65) | def __init__(self, approximate='none'):
    method forward (line 70) | def forward(self, X):
  class ComplexSiLU (line 74) | class ComplexSiLU(nn.Module):
    method __init__ (line 75) | def __init__(self):
    method forward (line 80) | def forward(self, X):
  class ComplexReLU (line 84) | class ComplexReLU(nn.Module):
    method __init__ (line 85) | def __init__(self):
    method forward (line 90) | def forward(self, X):
  class ComplexAvgPool3d (line 94) | class ComplexAvgPool3d(nn.Module):
    method __init__ (line 95) | def __init__(self, kernel_size, stride, padding):
    method forward (line 104) | def forward(self, X):
  class ComplexFlatten (line 108) | class ComplexFlatten(nn.Module):
    method __init__ (line 109) | def __init__(self, start_dim=1, end_dim=-1):
    method forward (line 114) | def forward(self, X):
  class NaiveComplexBatchNorm3d (line 118) | class NaiveComplexBatchNorm3d(nn.Module):
    method __init__ (line 119) | def __init__(
    method forward (line 135) | def forward(self, X):
  class NaiveComplexLayerNorm (line 139) | class NaiveComplexLayerNorm(nn.Module):
    method __init__ (line 140) | def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
    method forward (line 145) | def forward(self, X):
  class ComplexLinear (line 149) | class ComplexLinear(nn.Module):
    method __init__ (line 150) | def __init__(self, in_features, out_features, bias=True):
    method forward (line 155) | def forward(self, X):
  class ComplexMLP (line 159) | class ComplexMLP(nn.Module):
    method __init__ (line 160) | def __init__(self, in_features, hidden_features=None, out_features=Non...
    method forward (line 170) | def forward(self, x):
  class ComplexConv3d (line 178) | class ComplexConv3d(nn.Module):
    method __init__ (line 179) | def __init__(self, input_channels, num_channels, kernel_size, padding,...
    method forward (line 198) | def forward(self, X):
  class ComplexResidual3d (line 202) | class ComplexResidual3d(nn.Module):
    method __init__ (line 203) | def __init__(self, input_channels, num_channels, kernel_size, padding,...
    method forward (line 227) | def forward(self, X):
  class ComplexSegment (line 234) | class ComplexSegment(nn.Module):
    method __init__ (line 235) | def __init__(self, input_channels, seg_channels, seg_size):
    method forward (line 246) | def forward(self, X):
  class Complex2Real (line 253) | class Complex2Real(nn.Module):
    method __init__ (line 254) | def __init__(self):
    method forward (line 259) | def forward(self, X):
  class ComplexDotProductAttention (line 265) | class ComplexDotProductAttention(nn.Module):
    method __init__ (line 271) | def __init__(self, dropout, **kwargs):
    method forward (line 275) | def forward(self, queries, keys, values):
  class ComplexMultiHeadAttention (line 284) | class ComplexMultiHeadAttention(nn.Module):
    method __init__ (line 285) | def __init__(
    method forward (line 306) | def forward(self, queries, keys, values):
  class ComplexPositionalEncoding (line 316) | class ComplexPositionalEncoding(nn.Module):
    method __init__ (line 317) | def __init__(self, hidden_dim, dropout, max_len=10000):
    method forward (line 328) | def forward(self, X):
  class PositionWiseFFN (line 334) | class PositionWiseFFN(nn.Module):
    method __init__ (line 335) | def __init__(self, input_dim, hidden_dim, output_dim, **kwargs):
    method forward (line 341) | def forward(self, X):
  class ComplexAddNorm (line 346) | class ComplexAddNorm(nn.Module):
    method __init__ (line 347) | def __init__(self, normalized_shape, dropout, **kwargs):
    method forward (line 352) | def forward(self, X, Y):
  class ComplexEncoderBlock (line 357) | class ComplexEncoderBlock(nn.Module):
    method __init__ (line 358) | def __init__(
    method forward (line 380) | def forward(self, X):
  class ComplexTransformerEncoder (line 386) | class ComplexTransformerEncoder(nn.Module):
    method __init__ (line 387) | def __init__(
    method forward (line 423) | def forward(self, X, *args):

FILE: inference.py
  function gaussian (line 31) | def gaussian(window_size: int, tfdiff: float):
  function create_window (line 37) | def create_window(height: int, width: int):
  function eval_ssim (line 45) | def eval_ssim(pred, data, height, width, device):
  function cal_SNR_EEG (line 61) | def cal_SNR_EEG(predict, truth):
  function cal_SNR_MIMO (line 71) | def cal_SNR_MIMO(predict, truth):
  function save (line 84) | def save(out_dir, data, cond, batch, index=0):
  function save_mimo (line 93) | def save_mimo(out_dir, data, pred, cond, batch, index=0):
  function save_wifi (line 133) | def save_wifi(out_dir, data, pred, cond, batch, index=0):
  function save_fmcw (line 188) | def save_fmcw(out_dir, data, pred, cond, batch,index=0):
  function print_fid (line 251) | def print_fid(out_dir,data_dir,task_id):
  function main (line 265) | def main(args):

FILE: plots/code/Fig13(b)-Performance-of-channel-estimation-SNR.py
  function read_data_from_txt (line 8) | def read_data_from_txt(filename):

FILE: tfdiff/dataset.py
  function _nested_map (line 20) | def _nested_map(struct, map_fn):
  class WiFiDataset (line 30) | class WiFiDataset(torch.utils.data.Dataset):
    method __init__ (line 31) | def __init__(self, paths):
    method __len__ (line 37) | def __len__(self):
    method __getitem__ (line 40) | def __getitem__(self, idx):
  class FMCWDataset (line 52) | class FMCWDataset(torch.utils.data.Dataset):
    method __init__ (line 53) | def __init__(self, paths):
    method __len__ (line 59) | def __len__(self):
    method __getitem__ (line 62) | def __getitem__(self, idx):
  class MIMODataset (line 72) | class MIMODataset(torch.utils.data.Dataset):
    method __init__ (line 73) | def __init__(self, paths):
    method __len__ (line 79) | def __len__(self):
    method __getitem__ (line 82) | def __getitem__(self,idx):
  class EEGDataset (line 92) | class EEGDataset(torch.utils.data.Dataset):
    method __init__ (line 93) | def __init__(self, paths):
    method __len__ (line 99) | def __len__(self):
    method __getitem__ (line 102) | def __getitem__(self,idx):
  class Collator (line 112) | class Collator:
    method __init__ (line 113) | def __init__(self, params):
    method collate (line 116) | def collate(self, minibatch):
  function from_path (line 195) | def from_path(params, is_distributed=False):
  function from_path_inference (line 220) | def from_path_inference(params):

FILE: tfdiff/diffusion.py
  class SignalDiffusion (line 6) | class SignalDiffusion(nn.Module):
    method __init__ (line 7) | def __init__(self, params):
    method get_kernel (line 28) | def get_kernel(self, var_kernel):
    method get_noise_weights (line 34) | def get_noise_weights(self):
    method get_noise_weights_stats (line 50) | def get_noise_weights_stats(self):
    method get_noise_weights_div (line 58) | def get_noise_weights_div(self):
    method get_noise_weights_prod (line 69) | def get_noise_weights_prod(self):
    method degrade_fn (line 82) | def degrade_fn(self, x_0, t, task_id):
    method sampling (line 97) | def sampling(self, restore_fn, cond, device):
    method robust_sampling (line 121) | def robust_sampling(self, restore_fn, cond, device):
    method fast_sampling (line 144) | def fast_sampling(self, restore_fn, cond, device):
    method native_sampling (line 159) | def native_sampling(self, restore_fn, data, cond, device):
  class GaussianDiffusion (line 169) | class GaussianDiffusion(nn.Module):
    method __init__ (line 170) | def __init__(self, params):
    method degrade_fn (line 183) | def degrade_fn(self, x_0, t):
    method sampling (line 192) | def sampling(self, restore_fn, cond, device):
    method robust_sampling (line 206) | def robust_sampling(self, restore_fn, cond, device):
    method fast_sampling (line 220) | def fast_sampling(self, restore_fn, cond, device):
    method native_sampling (line 232) | def native_sampling(self, restore_fn, data, cond, device):

FILE: tfdiff/eeg_model.py
  function init_weight_norm (line 11) | def init_weight_norm(module):
  function init_weight_zero (line 18) | def init_weight_zero(module):
  function init_weight_xavier (line 25) | def init_weight_xavier(module):
  function modulate (line 33) | def modulate(x, shift, scale):
  class DiffusionEmbedding (line 37) | class DiffusionEmbedding(nn.Module):
    method __init__ (line 38) | def __init__(self, max_step, embed_dim=256, hidden_dim=256):
    method forward (line 50) | def forward(self, t):
    method _lerp_embedding (line 57) | def _lerp_embedding(self, t):
    method _build_embedding (line 64) | def _build_embedding(self, max_step, embed_dim):
  class MLPConditionEmbedding (line 74) | class MLPConditionEmbedding(nn.Module):
    method __init__ (line 75) | def __init__(self, cond_dim, hidden_dim=256):
    method forward (line 86) | def forward(self, c):
  class PositionEmbedding (line 89) | class PositionEmbedding(nn.Module):
    method __init__ (line 90) | def __init__(self, max_len, input_dim, hidden_dim):
    method forward (line 97) | def forward(self, x):
    method _build_embedding (line 101) | def _build_embedding(self, max_len, hidden_dim):
  class CDiTBlock (line 110) | class CDiTBlock(nn.Module):
    method __init__ (line 111) | def __init__(self, hidden_dim, num_heads, dropout, mlp_ratio=4.0, **bl...
    method forward (line 130) | def forward(self, x, c):
  class FinalLayer (line 144) | class FinalLayer(nn.Module):
    method __init__ (line 145) | def __init__(self, hidden_dim, out_dim):
    method forward (line 156) | def forward(self, x, c):
  class tfdiff_eeg (line 163) | class tfdiff_eeg(nn.Module):
    method __init__ (line 164) | def __init__(self, params):
    method forward (line 186) | def forward(self, x, t, c):

FILE: tfdiff/fmcw_model.py
  function init_weight_norm (line 11) | def init_weight_norm(module):
  function init_weight_zero (line 18) | def init_weight_zero(module):
  function init_weight_xavier (line 25) | def init_weight_xavier(module):
  function modulate (line 33) | def modulate(x, shift, scale):
  class DiffusionEmbedding (line 37) | class DiffusionEmbedding(nn.Module):
    method __init__ (line 38) | def __init__(self, max_step, embed_dim=256, hidden_dim=256):
    method forward (line 50) | def forward(self, t):
    method _lerp_embedding (line 57) | def _lerp_embedding(self, t):
    method _build_embedding (line 64) | def _build_embedding(self, max_step, embed_dim):
  class MLPConditionEmbedding (line 74) | class MLPConditionEmbedding(nn.Module):
    method __init__ (line 75) | def __init__(self, cond_dim, hidden_dim=256):
    method forward (line 86) | def forward(self, c):
  class PositionEmbedding (line 90) | class PositionEmbedding(nn.Module):
    method __init__ (line 91) | def __init__(self, max_len, input_dim, hidden_dim):
    method forward (line 98) | def forward(self, x):
    method _build_embedding (line 102) | def _build_embedding(self, max_len, hidden_dim):
  class DiA (line 111) | class DiA(nn.Module):
    method __init__ (line 112) | def __init__(self, hidden_dim, num_heads, dropout, mlp_ratio=4.0, **bl...
    method forward (line 133) | def forward(self, x, c):
  class FinalLayer (line 154) | class FinalLayer(nn.Module):
    method __init__ (line 155) | def __init__(self, hidden_dim, out_dim):
    method forward (line 166) | def forward(self, x, c):
  class tfdiff_fmcw (line 320) | class tfdiff_fmcw(nn.Module):
    method __init__ (line 321) | def __init__(self, params):
    method forward (line 343) | def forward(self, x, t, c):

FILE: tfdiff/learner.py
  class tfdiffLoss (line 11) | class tfdiffLoss(nn.Module):
    method __init__ (line 12) | def __init__(self, w=0.1):
    method forward (line 16) | def forward(self, target, est, target_noise=None, est_noise=None):
    method complex_mse_loss (line 24) | def complex_mse_loss(self, target, est):
  class tfdiffLearner (line 30) | class tfdiffLearner:
    method __init__ (line 31) | def __init__(self, log_dir, model_dir, model, dataset, optimizer, para...
    method state_dict (line 58) | def state_dict(self):
    method load_state_dict (line 70) | def load_state_dict(self, state_dict):
    method save_to_checkpoint (line 78) | def save_to_checkpoint(self, filename='weights'):
    method restore_from_checkpoint (line 90) | def restore_from_checkpoint(self, filename='weights'):
    method train (line 98) | def train(self, max_iter=None):
    method train_iter (line 121) | def train_iter(self, features):
    method _write_summary (line 140) | def _write_summary(self, iter, features, loss):

FILE: tfdiff/mimo_model.py
  function init_weight_norm (line 11) | def init_weight_norm(module):
  function init_weight_zero (line 18) | def init_weight_zero(module):
  function init_weight_xavier (line 25) | def init_weight_xavier(module):
  function modulate (line 33) | def modulate(x, shift, scale):
  class DiffusionEmbedding (line 37) | class DiffusionEmbedding(nn.Module):
    method __init__ (line 38) | def __init__(self, max_step, embed_dim=256, hidden_dim=256):
    method forward (line 50) | def forward(self, t):
    method _lerp_embedding (line 57) | def _lerp_embedding(self, t):
    method _build_embedding (line 64) | def _build_embedding(self, max_step, embed_dim):
  class MLPConditionEmbedding (line 74) | class MLPConditionEmbedding(nn.Module):
    method __init__ (line 75) | def __init__(self, cond_dim, hidden_dim=256):
    method forward (line 86) | def forward(self, c):
  class PositionEmbedding (line 90) | class PositionEmbedding(nn.Module):
    method __init__ (line 91) | def __init__(self, max_len, input_dim, hidden_dim):
    method forward (line 98) | def forward(self, x):
    method _build_embedding (line 102) | def _build_embedding(self, max_len, hidden_dim):
  class DiA (line 111) | class DiA(nn.Module):
    method __init__ (line 112) | def __init__(self, hidden_dim, num_heads, dropout, mlp_ratio=4.0, **bl...
    method forward (line 139) | def forward(self, x, t, c):
  class FinalLayer (line 162) | class FinalLayer(nn.Module):
    method __init__ (line 163) | def __init__(self, hidden_dim, out_dim):
    method forward (line 174) | def forward(self, x, t):
  class SpatialDiffusion (line 181) | class SpatialDiffusion(nn.Module):
    method __init__ (line 193) | def __init__(self, params):
    method forward (line 226) | def forward(self, x, t, c):
  class TimeFrequencyDiffusion (line 236) | class TimeFrequencyDiffusion(nn.Module):
    method __init__ (line 248) | def __init__(self, params):
    method forward (line 274) | def forward(self, x, t, c):
  class tfdiff_mimo (line 285) | class tfdiff_mimo(nn.Module):
    method __init__ (line 297) | def __init__(self, params):
    method forward (line 309) | def forward(self, x, t, c):

FILE: tfdiff/params.py
  class AttrDict (line 4) | class AttrDict(dict):
    method __init__ (line 5) | def __init__(self, *args, **kwargs):
    method override (line 9) | def override(self, attrs):

FILE: tfdiff/wifi_model.py
  function init_weight_norm (line 11) | def init_weight_norm(module):
  function init_weight_zero (line 18) | def init_weight_zero(module):
  function init_weight_xavier (line 25) | def init_weight_xavier(module):
  function modulate (line 33) | def modulate(x, shift, scale):
  class DiffusionEmbedding (line 37) | class DiffusionEmbedding(nn.Module):
    method __init__ (line 38) | def __init__(self, max_step, embed_dim=256, hidden_dim=256):
    method forward (line 50) | def forward(self, t):
    method _lerp_embedding (line 57) | def _lerp_embedding(self, t):
    method _build_embedding (line 64) | def _build_embedding(self, max_step, embed_dim):
  class MLPConditionEmbedding (line 74) | class MLPConditionEmbedding(nn.Module):
    method __init__ (line 75) | def __init__(self, cond_dim, hidden_dim=256):
    method forward (line 86) | def forward(self, c):
  class PositionEmbedding (line 90) | class PositionEmbedding(nn.Module):
    method __init__ (line 91) | def __init__(self, max_len, input_dim, hidden_dim):
    method forward (line 98) | def forward(self, x):
    method _build_embedding (line 102) | def _build_embedding(self, max_len, hidden_dim):
  class DiA (line 111) | class DiA(nn.Module):
    method __init__ (line 112) | def __init__(self, hidden_dim, num_heads, dropout, mlp_ratio=4.0, **bl...
    method forward (line 133) | def forward(self, x, c):
  class FinalLayer (line 154) | class FinalLayer(nn.Module):
    method __init__ (line 155) | def __init__(self, hidden_dim, out_dim):
    method forward (line 166) | def forward(self, x, c):
  class tfdiff_WiFi (line 320) | class tfdiff_WiFi(nn.Module):
    method __init__ (line 321) | def __init__(self, params):
    method forward (line 343) | def forward(self, x, t, c):

FILE: train.py
  function _get_free_port (line 18) | def _get_free_port():
  function _train_impl (line 23) | def _train_impl(replica_id, model, dataset, params):
  function train (line 31) | def train(params):
  function train_distributed (line 42) | def train_distributed(replica_id, replica_count, port, params):
  function main (line 64) | def main(args):