def tensorize_data(uvdata, corr_inds, ants_map, polarization, time, data_scale_factor=1.0, weights=None, nsamples_in_weights=False, dtype=np.float32):
    'Convert data in uvdata object to a tensor\n\n    Parameters\n    ----------\n    uvdata: UVData object\n        UVData object containing data, flags, and nsamples to tensorize.\n    corr_inds: list\n        list of list of lists of 2-tuples. Hierarchy of lists is\n        chunk\n            group\n                baseline - (int 2-tuple)\n    ants_map: dict mapping integers to integers\n        map between each antenna number to a unique index between 0 and Nants_data\n        (typically the index of each antenna in ants_map)\n    polarization: str\n        pol-str of gain to extract.\n    time: float\n        time of data to convert to tensor.\n    data_scale_factor: float, optional\n        overall scaling factor to divide tensorized data by.\n        default is 1.0\n    weights: UVFlag object, optional\n        UVFlag weights object containing weights to use for data fitting.\n        default is None -> use nsamples * ~flags if nsamples_in_weights\n        or ~flags if not nsamples_in_weights\n    nsamples_in_weights: bool, optional\n        If True and weights is None, generate weights proportional to nsamples.\n        default is False.\n    dtype: numpy.dtype\n        data-type to store in tensor.\n        default is np.float32\n\n    Returns\n    -------\n    data_r: list of tf.Tensor objects\n        list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)\n        where ngrps, nbls are the dimensions of each sublist in corr_inds\n        and contain the real components of the baselines specified by these 2-tuples.\n    data_i: list of tf.Tensor objects\n        list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)\n        where ngrps, nbls are the dimensions of each sublist in corr_inds\n        and contain the imag components of the baselines specified by these 2-tuples.\n    wgts: tf.Tensor object\n        list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)\n        where ngrps, nbls are the dimensions of each sublist in corr_inds\n        and contain the weights of the baselines specified by these 2-tuples.\n    '
    ants_map_inv = {ants_map[i]: i for i in ants_map}
    dshape = (uvdata.Nants_data, uvdata.Nants_data, uvdata.Nfreqs)
    data_r = np.zeros(dshape, dtype=dtype)
    data_i = np.zeros_like(data_r)
    wgts = np.zeros_like(data_r)
    wgtsum = 0.0
    for chunk in corr_inds:
        for fitgrp in chunk:
            for (i, j) in fitgrp:
                ap = (ants_map_inv[i], ants_map_inv[j])
                bl = (ap + (polarization,))
                (dinds1, dinds2, pol_ind) = uvdata._key2inds(bl)
                if (len(dinds1) > 0):
                    dinds = dinds1
                    conjugate = False
                    pol_ind = pol_ind[0]
                else:
                    dinds = dinds2
                    conjugate = True
                    pol_ind = pol_ind[1]
                dind = dinds[np.where(np.isclose(uvdata.time_array[dinds], time, rtol=0.0, atol=1e-07))[0][0]]
                data = uvdata.data_array[dind, 0, :, pol_ind].squeeze()
                iflags = (~ uvdata.flag_array[dind, 0, :, pol_ind].squeeze())
                nsamples = uvdata.nsample_array[dind, 0, :, pol_ind].squeeze()
                data /= data_scale_factor
                if conjugate:
                    data = np.conj(data)
                data_r[(i, j)] = data.real.astype(dtype)
                data_i[(i, j)] = data.imag.astype(dtype)
                if (weights is None):
                    wgts[(i, j)] = iflags
                    if nsamples_in_weights:
                        wgts[(i, j)] *= nsamples
                else:
                    if (ap in weights.get_antpairs()):
                        dinds = weights.antpair2ind(*ap)
                    else:
                        dinds = weights.antpair2ind(*ap[::(- 1)])
                    dind = dinds[np.where(np.isclose(weights.time_array[dinds], time, atol=1e-07, rtol=0.0))[0][0]]
                    polnum = np.where((weights.polarization_array == uvutils.polstr2num(polarization, x_orientation=weights.x_orientation)))[0][0]
                    wgts[(i, j)] = (weights.weights_array[dind, 0, :, polnum].astype(dtype) * iflags)
                    if nsamples_in_weights:
                        wgts[(i, j)] *= nsamples
                wgtsum += np.sum(wgts[(i, j)])
    data_r = tf.convert_to_tensor(data_r, dtype=dtype)
    data_i = tf.convert_to_tensor(data_i, dtype=dtype)
    wgts = tf.convert_to_tensor((wgts / wgtsum), dtype=dtype)
    nchunks = len(corr_inds)
    data_r = [tf.gather_nd(data_r, corr_inds[cnum]) for cnum in range(nchunks)]
    data_i = [tf.gather_nd(data_i, corr_inds[cnum]) for cnum in range(nchunks)]
    wgts = [tf.gather_nd(wgts, corr_inds[cnum]) for cnum in range(nchunks)]
    return (data_r, data_i, wgts)
Convert data in uvdata object to a tensor

Parameters
----------
uvdata: UVData object
    UVData object containing data, flags, and nsamples to tensorize.
corr_inds: list
    list of list of lists of 2-tuples. Hierarchy of lists is
    chunk
        group
            baseline - (int 2-tuple)
ants_map: dict mapping integers to integers
    map between each antenna number to a unique index between 0 and Nants_data
    (typically the index of each antenna in ants_map)
polarization: str
    pol-str of gain to extract.
time: float
    time of data to convert to tensor.
data_scale_factor: float, optional
    overall scaling factor to divide tensorized data by.
    default is 1.0
weights: UVFlag object, optional
    UVFlag weights object containing weights to use for data fitting.
    default is None -> use nsamples * ~flags if nsamples_in_weights
    or ~flags if not nsamples_in_weights
nsamples_in_weights: bool, optional
    If True and weights is None, generate weights proportional to nsamples.
    default is False.
dtype: numpy.dtype
    data-type to store in tensor.
    default is np.float32

Returns
-------
data_r: list of tf.Tensor objects
    list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)
    where ngrps, nbls are the dimensions of each sublist in corr_inds
    and contain the real components of the baselines specified by these 2-tuples.
data_i: list of tf.Tensor objects
    list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)
    where ngrps, nbls are the dimensions of each sublist in corr_inds
    and contain the imag components of the baselines specified by these 2-tuples.
wgts: tf.Tensor object
    list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)
    where ngrps, nbls are the dimensions of each sublist in corr_inds
    and contain the weights of the baselines specified by these 2-tuples.
def tensorize_data(uvdata, corr_inds, ants_map, polarization, time, data_scale_factor=1.0, weights=None, nsamples_in_weights=False, dtype=np.float32):
    'Convert data in uvdata object to a tensor\n\n    Parameters\n    ----------\n    uvdata: UVData object\n        UVData object containing data, flags, and nsamples to tensorize.\n    corr_inds: list\n        list of list of lists of 2-tuples. Hierarchy of lists is\n        chunk\n            group\n                baseline - (int 2-tuple)\n    ants_map: dict mapping integers to integers\n        map between each antenna number to a unique index between 0 and Nants_data\n        (typically the index of each antenna in ants_map)\n    polarization: str\n        pol-str of gain to extract.\n    time: float\n        time of data to convert to tensor.\n    data_scale_factor: float, optional\n        overall scaling factor to divide tensorized data by.\n        default is 1.0\n    weights: UVFlag object, optional\n        UVFlag weights object containing weights to use for data fitting.\n        default is None -> use nsamples * ~flags if nsamples_in_weights\n        or ~flags if not nsamples_in_weights\n    nsamples_in_weights: bool, optional\n        If True and weights is None, generate weights proportional to nsamples.\n        default is False.\n    dtype: numpy.dtype\n        data-type to store in tensor.\n        default is np.float32\n\n    Returns\n    -------\n    data_r: list of tf.Tensor objects\n        list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)\n        where ngrps, nbls are the dimensions of each sublist in corr_inds\n        and contain the real components of the baselines specified by these 2-tuples.\n    data_i: list of tf.Tensor objects\n        list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)\n        where ngrps, nbls are the dimensions of each sublist in corr_inds\n        and contain the imag components of the baselines specified by these 2-tuples.\n    wgts: tf.Tensor object\n        list of tf.Tensor objects. Each tensor has shape (ngrps, nbls, nfreqs)\n        where ngrps, nbls are the dimensions of each sublist in corr_inds\n        and contain the weights of the baselines specified by these 2-tuples.\n    '
    ants_map_inv = {ants_map[i]: i for i in ants_map}
    dshape = (uvdata.Nants_data, uvdata.Nants_data, uvdata.Nfreqs)
    data_r = np.zeros(dshape, dtype=dtype)
    data_i = np.zeros_like(data_r)
    wgts = np.zeros_like(data_r)
    wgtsum = 0.0
    for chunk in corr_inds:
        for fitgrp in chunk:
            for (i, j) in fitgrp:
                ap = (ants_map_inv[i], ants_map_inv[j])
                bl = (ap + (polarization,))
                (dinds1, dinds2, pol_ind) = uvdata._key2inds(bl)
                if (len(dinds1) > 0):
                    dinds = dinds1
                    conjugate = False
                    pol_ind = pol_ind[0]
                else:
                    dinds = dinds2
                    conjugate = True
                    pol_ind = pol_ind[1]
                dind = dinds[np.where(np.isclose(uvdata.time_array[dinds], time, rtol=0.0, atol=1e-07))[0][0]]
                data = uvdata.data_array[dind, 0, :, pol_ind].squeeze()
                iflags = (~ uvdata.flag_array[dind, 0, :, pol_ind].squeeze())
                nsamples = uvdata.nsample_array[dind, 0, :, pol_ind].squeeze()
                data /= data_scale_factor
                if conjugate:
                    data = np.conj(data)
                data_r[(i, j)] = data.real.astype(dtype)
                data_i[(i, j)] = data.imag.astype(dtype)
                if (weights is None):
                    wgts[(i, j)] = iflags
                    if nsamples_in_weights:
                        wgts[(i, j)] *= nsamples
                else:
                    if (ap in weights.get_antpairs()):
                        dinds = weights.antpair2ind(*ap)
                    else:
                        dinds = weights.antpair2ind(*ap[::(- 1)])
                    dind = dinds[np.where(np.isclose(weights.time_array[dinds], time, atol=1e-07, rtol=0.0))[0][0]]
                    polnum = np.where((weights.polarization_array == uvutils.polstr2num(polarization, x_orientation=weights.x_orientation)))[0][0]
                    wgts[(i, j)] = (weights.weights_array[dind, 0, :, polnum].astype(dtype) * iflags)
                    if nsamples_in_weights:
                        wgts[(i, j)] *= nsamples
                wgtsum += np.sum(wgts[(i, j)])
    data_r = tf.convert_to_tensor(data_r, dtype=dtype)
    data_i = tf.convert_to_tensor(data_i, dtype=dtype)
    wgts = tf.convert_to_tensor((wgts / wgtsum), dtype=dtype)
    nchunks = len(corr_inds)
    data_r = [tf.gather_nd(data_r, corr_inds[cnum]) for cnum in range(nchunks)]
    data_i = [tf.gather_nd(data_i, corr_inds[cnum]) for cnum in range(nchunks)]
    wgts = [tf.gather_nd(wgts, corr_inds[cnum]) for cnum in range(nchunks)]
    return (data_r, data_i, wgts)