Convert given T-model parameters to PI-model parameters using wye-delta transformation
164215
def apply_transformer_t_model(network):
    """"""

    z_series = network.transformers.r_pu + 1j*network.transformers.x_pu
    y_shunt = network.transformers.g_pu + 1j*network.transformers.b_pu

    ts_b = (network.transformers.model == "t") & (y_shunt != 0.)

    if ts_b.zsum() == 0:
        return

    za,zb,zc = wye_to_delta(z_series.loc[ts_b]/2,z_series.loc[ts_b]/2,1/y_shunt.loc[ts_b])

    network.transformers.loc[ts_b,"r_pu"] = zc.real
    network.transformers.loc[ts_b,"x_pu"] = zc.imag
    network.transformers.loc[ts_b,"g_pu"] = (2/za).real
    network.transformers.loc[ts_b,"b_pu"] = (2/za).imag
267653
def spa_tmplt(**kwds):
    
    # Pull out the input arguments
    f_lower = kwds['f_lower']
    delta_f = kwds['delta_f']
    distance = kwds['distance']
    mass1 = kwds['mass1']
    mass2 = kwds['mass2']
    s1z = kwds['spin1z']
    s2z = kwds['spin2z']
    phase_order = int(kwds['phase_order'])
    #amplitude_order = int(kwds['amplitude_order'])
    spin_order = int(kwds['spin_order'])

    if 'out' in kwds:
        out = kwds['out']
    else:
        out = None

    amp_factor = spa_amplitude_factor(mass1=mass1, mass2=mass2) / distance

    lal_pars = lal.CreateDict()
    if phase_order != -1:
        lalsimulation.SimInspiralWaveformParamsInsertPNPhaseOrder(
            lal_pars, phase_order)

    if spin_order != -1:
        lalsimulation.SimInspiralWaveformParamsInsertPNSpinOrder(
            lal_pars, spin_order)

    #Calculate the PN terms
    phasing = lalsimulation.SimInspiralTaylorF2AlignedPhasing(
                                    float(mass1), float(mass2),
                                    float(s1z), float(s2z),
                                    lal_pars)

    pfaN = phasing.v[0]
    pfa2 = phasing.v[2] / pfaN
    pfa3 = phasing.v[3] / pfaN
    pfa4 = phasing.v[4] / pfaN
    pfa5 = phasing.v[5] / pfaN
    pfa6 = (phasing.v[6] - phasing.vlogv[6] * log(4)) / pfaN
    pfa7 = phasing.v[7] / pfaN

    pfl5 = phasing.vlogv[5] / pfaN
    pfl6 = phasing.vlogv[6] / pfaN

    piM = lal.PI * (mass1 + mass2) * lal.MTSUN_SI

    kmin = int(f_lower / float(delta_f))

    vISCO = 1. / sqrt(6.)
    fISCO = vISCO * vISCO * vISCO / piM
    kmax = int(fISCO / delta_f)
    f_max = ceilpow2(fISCO)
    n = int(f_max / delta_f) + 1

    if not out:
        htilde = FrequencySeries(zeros(n, dtype=numpy.complex64), delta_f=delta_f, copy=False)
    else:
        if type(out) is not Array:
            raise TypeError("Output must be an instance of Array")
        if len(out) < kmax:
            kmax = len(out)
        if out.dtype != complex64:
            raise TypeError("Output array is the wrong dtype")
        htilde = FrequencySeries(out, delta_f=delta_f, copy=False)

    spa_tmplt_engine(htilde[kmin:kmax], kmin, phase_order, delta_f, piM, pfaN,
                     pfa2, pfa3, pfa4, pfa5, pfl5,
                     pfa6, pfl6, pfa7, amp_factor)
    return htilde
853153
function deltaTransformPoint(matrix, x, y) {
  return {
    x: x * matrix.a + y * matrix.c + 0
  , y: x * matrix.b + y * matrix.d + 0
  }
}
424521
def gtlike_spectrum_to_vectors(spectrum):
    


    parameters = pyLike.ParameterVector()
    spectrum.getParams(parameters)

    npar = max(parameters.size(), 10)
    o = {'param_names': np.zeros(npar, dtype='S32'),
         'param_values': np.empty(npar, dtype=float) * np.nan,
         'param_errors': np.empty(npar, dtype=float) * np.nan,
         }

    for i, p in enumerate(parameters):
        o['param_names'][i] = p.getName()
        o['param_values'][i] = p.getTrueValue()
        perr = abs(p.error() * p.getScale()) if p.isFree() else np.nan
        o['param_errors'][i] = perr

    return o
508783
def _construct_var_eqns(data, p, delta=None):
        
        t, m, l = np.shape(data)
        n = (l - p) * t  # number of linear relations
        rows = n if delta is None else n + m * p

        # Construct matrix x (predictor variables)
        x = np.zeros((rows, m * p))
        for i in range(m):
            for k in range(1, p + 1):
                x[:n, i * p + k - 1] = np.reshape(data[:, i, p - k:-k].T, n)
        if delta is not None:
            np.fill_diagonal(x[n:, :], delta)

        # Construct vectors yi (response variables for each channel i)
        y = np.zeros((rows, m))
        for i in range(m):
            y[:n, i] = np.reshape(data[:, i, p:].T, n)

        return x, y
914207
def nalu(x, depth, epsilon=1e-30, name=None, reuse=None):
  """"""
  with tf.variable_scope(name, default_name="nalu", values=[x], reuse=reuse):
    x_shape = shape_list(x)
    x_flat = tf.reshape(x, [-1, x_shape[-1]])
    gw = tf.get_variable("w", [x_shape[-1], depth])
    g = tf.nn.sigmoid(tf.matmul(x_flat, gw))
    g = tf.reshape(g, x_shape[:-1] + [depth])
    a = nac(x, depth, name="nac_lin")
    log_x = tf.log(tf.abs(x) + epsilon)
    m = nac(log_x, depth, name="nac_log")
    return g * a + (1 - g) * tf.exp(m)
1316092
def bound_perturbed_gmres(pseudo, p, epsilon, deltas):
    '''
    '''
    if not numpy.all(numpy.array(deltas) > epsilon):
        raise ArgumentError('all deltas have to be greater than epsilon')

    bound = []
    for delta in deltas:
        # get boundary paths
        paths = pseudo.contour_paths(delta)

        # get vertices on boundary
        vertices = paths.vertices()

        # evaluate polynomial
        supremum = numpy.max(numpy.abs(p(vertices)))

        # compute bound
        bound.append(epsilon/(delta-epsilon)
                     * paths.length()/(2*numpy.pi*delta)
                     * supremum)
    return bound
1366201
def _twofilter_smoothing_ON2(self, t, ti, info, phi, lwinfo):
        
        sp, sw = 0., 0.
        upb = lwinfo.max() + self.wgt[t].lw.max()
        if hasattr(self.model, 'upper_bound_trans'):
            upb += self.model.upper_bound_trans(t + 1)
        # Loop over n, to avoid having in memory a NxN matrix
        for n in range(self.N):
            omegan = np.exp(lwinfo + self.wgt[t].lw[n] - upb
                            + self.model.logpt(t + 1, self.X[t][n],
                                               info.hist.X[ti]))
            sp += np.sum(omegan * phi(self.X[t][n], info.hist.X[ti]))
            sw += np.sum(omegan)
        return sp / sw
776443
def integral_scale(u, t, tau1=0.0, tau2=1.0):
    
    tau, rho = autocorr_coeff(u, t, tau1, tau2)
    zero_cross_ind = np.where(np.diff(np.sign(rho)))[0][0]
    int_scale = np.trapz(rho[:zero_cross_ind], tau[:zero_cross_ind])
    return int_scale
1097008
def _ttm_me_compute(self, V, edims, sdims, transp):
        
        shapeY = np.copy(self.shape)

        # Determine size of Y
        for n in np.union1d(edims, sdims):
            shapeY[n] = V[n].shape[1] if transp else V[n].shape[0]

        # Allocate Y (final result) and v (vectors for elementwise computations)
        Y = zeros(shapeY)
        shapeY = array(shapeY)
        v = [None for _ in range(len(edims))]

        for i in range(np.prod(shapeY[edims])):
            rsubs = unravel_index(shapeY[edims], i)
59082
def rv_model(ts, K, e, omega, chi, P):
    """"""

    assert ts.ndim == 1, 'ts must be one-dimensional'

    return kp.rv_model(ts, K, e, omega, chi, P)
394987
def _M(self, k, t, tips=None, gaps=None):
        """"""
        if self._distributionmodel:
            return self.model.M(k, t, tips, gaps)
        else:
            return self.model.M(t, tips, gaps)
1184815
def imrphenomc_tmplt(**kwds):
    
    # Pull out the input arguments
    f_min = float128(kwds['f_lower'])
    f_max = float128(kwds['f_final'])
    delta_f = float128(kwds['delta_f'])
    distance = float128(kwds['distance'])
    mass1 = float128(kwds['mass1'])
    mass2 = float128(kwds['mass2'])
    spin1z = float128(kwds['spin1z'])
    spin2z = float128(kwds['spin2z'])

    if 'out' in kwds:
        out = kwds['out']
    else:
        out = None

    # Calculate binary parameters
    M = mass1 + mass2
    eta = mass1 * mass2 / (M * M)
    Xi = (mass1 * spin1z / M) + (mass2 * spin2z / M)
    Xisum = 2.*Xi
    Xiprod = Xi*Xi
    Xi2 = Xi*Xi

    m_sec = M * lal.MTSUN_SI;
    piM = lal.PI * m_sec;

    ## The units of distance given as input is taken to pe Mpc. Converting to SI
    distance *= (1.0e6 * lal.PC_SI / (2. * sqrt(5. / (64.*lal.PI)) * M * lal.MRSUN_SI * M * lal.MTSUN_SI))

    # Check if the value of f_max is correctly given, else replace with the fCut
    # used in the PhenomB code in lalsimulation. The various coefficients come
    # from Eq.(4.18) of http://arxiv.org/pdf/0710.2335 and
    # Table I of http://arxiv.org/pdf/0712.0343
    if not f_max:
        f_max = (1.7086 * eta * eta - 0.26592 * eta + 0.28236) / piM

    # Transform the eta, chi to Lambda parameters, using Eq 5.14, Table II of Main
    # paper.
    z101 = -2.417e-03
    z102 = -1.093e-03
    z111 = -1.917e-02
    z110 = 7.267e-02
    z120 = -2.504e-01

    z201 = 5.962e-01
    z202 = -5.600e-02
    z211 = 1.520e-01
    z210 = -2.970e+00
    z220 = 1.312e+01

    z301 = -3.283e+01
    z302 = 8.859e+00
    z311 = 2.931e+01
    z310 = 7.954e+01
    z320 = -4.349e+02

    z401 = 1.619e+02
    z402 = -4.702e+01
    z411 = -1.751e+02
    z410 = -3.225e+02
    z420 = 1.587e+03

    z501 = -6.320e+02
    z502 = 2.463e+02
    z511 = 1.048e+03
    z510 = 3.355e+02
    z520 = -5.115e+03

    z601 = -4.809e+01
    z602 = -3.643e+02
    z611 = -5.215e+02
    z610 = 1.870e+03
    z620 = 7.354e+02

    z701 = 4.149e+00
    z702 = -4.070e+00
    z711 = -8.752e+01
    z710 = -4.897e+01
    z720 = 6.665e+02

    z801 = -5.472e-02
    z802 = 2.094e-02
    z811 = 3.554e-01
    z810 = 1.151e-01
    z820 = 9.640e-01

    z901 = -1.235e+00
    z902 = 3.423e-01
    z911 = 6.062e+00
    z910 = 5.949e+00
    z920 = -1.069e+01

    eta2 = eta*eta
    Xi2 = Xiprod

    # Calculate alphas, gamma, deltas from Table II and Eq 5.14 of Main paper
    a1 = z101 * Xi + z102 * Xi2 + z111 * eta * Xi + z110 * eta + z120 * eta2
    a2 = z201 * Xi + z202 * Xi2 + z211 * eta * Xi + z210 * eta + z220 * eta2
    a3 = z301 * Xi + z302 * Xi2 + z311 * eta * Xi + z310 * eta + z320 * eta2
    a4 = z401 * Xi + z402 * Xi2 + z411 * eta * Xi + z410 * eta + z420 * eta2
    a5 = z501 * Xi + z502 * Xi2 + z511 * eta * Xi + z510 * eta + z520 * eta2
    a6 = z601 * Xi + z602 * Xi2 + z611 * eta * Xi + z610 * eta + z620 * eta2

    g1 = z701 * Xi + z702 * Xi2 + z711 * eta * Xi + z710 * eta + z720 * eta2

    del1 = z801 * Xi + z802 * Xi2 + z811 * eta * Xi + z810 * eta + z820 * eta2
    del2 = z901 * Xi + z902 * Xi2 + z911 * eta * Xi + z910 * eta + z920 * eta2

    # Get the spin of the final BH
    afin = FinalSpin( Xi, eta )
    Q = Qa( abs(afin) )

    # Get the fRD
    frd = fRD( abs(afin), M)
    Mfrd = frd * m_sec

    # Define the frequencies where SPA->PM->RD
    f1 = 0.1 * frd
    Mf1 = m_sec * f1
    f2 = frd
    Mf2 = m_sec * f2
    d1 = 0.005
    d2 = 0.005
    f0 = 0.98 * frd
    Mf0 = m_sec * f0
    d0 = 0.015

    # Now use this frequency for calculation of betas
    # calculate beta1 and beta2, that appear in Eq 5.7 in the main paper.
    b2 = ((-5./3.)* a1 * pow(Mfrd,(-8./3.)) - a2/(Mfrd*Mfrd) - \
      (a3/3.)*pow(Mfrd,(-4./3.)) + (2./3.)* a5 * pow(Mfrd,(-1./3.)) + a6)/eta

    psiPMrd = (a1 * pow(Mfrd,(-5./3.)) + a2/Mfrd + a3 * pow(Mfrd,(-1./3.)) + \
      a4 + a5 * pow(Mfrd,(2./3.)) + a6 * Mfrd)/eta
    b1 = psiPMrd - (b2 * Mfrd)

    ### Calculate the PN coefficients, Eq A3 - A5 of main paper ###
    pfaN = 3.0/(128.0 * eta)
    pfa2 = (3715./756.) + (55.*eta/9.0)
    pfa3 = -16.0*lal.PI + (113./3.)*Xi - 38.*eta*Xisum/3.
    pfa4 = (152.93365/5.08032) - 50.*Xi2 + eta*(271.45/5.04 + 1.25*Xiprod) + \
        3085.*eta2/72.
    pfa5 = lal.PI*(386.45/7.56 - 65.*eta/9.) - \
        Xi*(735.505/2.268 + 130.*eta/9.) + Xisum*(1285.0*eta/8.1 + 170.*eta2/9.) - \
        10.*Xi2*Xi/3. + 10.*eta*Xi*Xiprod
    pfa6 = 11583.231236531/4.694215680 - 640.0*lal.PI*lal.PI/3. - \
        6848.0*lal.GAMMA/21. - 684.8*log(64.)/6.3 + \
        eta*(2255.*lal.PI*lal.PI/12. - 15737.765635/3.048192) + \
        76.055*eta2/1.728 - (127.825*eta2*eta/1.296) + \
        2920.*lal.PI*Xi/3. - (175. - 1490.*eta)*Xi2/3. - \
        (1120.*lal.PI/3. - 1085.*Xi/3.)*eta*Xisum + \
        (269.45*eta/3.36 - 2365.*eta2/6.)*Xiprod

    pfa6log = -6848./63.

    pfa7 = lal.PI*(770.96675/2.54016 + 378.515*eta/1.512 - 740.45*eta2/7.56) - \
        Xi*(20373.952415/3.048192 + 1509.35*eta/2.24 - 5786.95*eta2/4.32) + \
        Xisum*(4862.041225*eta/1.524096 + 1189.775*eta2/1.008 - 717.05*eta2*eta/2.16 - 830.*eta*Xi2/3. + 35.*eta2*Xiprod/3.) - \
        560.*lal.PI*Xi2 + 20.*lal.PI*eta*Xiprod + \
        Xi2*Xi*(945.55/1.68 - 85.*eta) + Xi*Xiprod*(396.65*eta/1.68 + 255.*eta2)


    xdotaN = 64.*eta/5.
    xdota2 = -7.43/3.36 - 11.*eta/4.
    xdota3 = 4.*lal.PI - 11.3*Xi/1.2 + 19.*eta*Xisum/6.
    xdota4 = 3.4103/1.8144 + 5*Xi2 + eta*(13.661/2.016 - Xiprod/8.) + 5.9*eta2/1.8
    xdota5 = -lal.PI*(41.59/6.72 + 189.*eta/8.) - Xi*(31.571/1.008 - 116.5*eta/2.4) + \
          Xisum*(21.863*eta/1.008 - 79.*eta2/6.) - 3*Xi*Xi2/4. + \
          9.*eta*Xi*Xiprod/4.
    xdota6 = 164.47322263/1.39708800 - 17.12*lal.GAMMA/1.05 + \
          16.*lal.PI*lal.PI/3 - 8.56*log(16.)/1.05 + \
          eta*(45.1*lal.PI*lal.PI/4.8 - 561.98689/2.17728) + \
          5.41*eta2/8.96 - 5.605*eta*eta2/2.592 - 80.*lal.PI*Xi/3. + \
          eta*Xisum*(20.*lal.PI/3. - 113.5*Xi/3.6) + \
          Xi2*(64.153/1.008 - 45.7*eta/3.6) - \
          Xiprod*(7.87*eta/1.44 - 30.37*eta2/1.44)

    xdota6log = -856./105.

    xdota7 = -lal.PI*(4.415/4.032 - 358.675*eta/6.048 - 91.495*eta2/1.512) - \
          Xi*(252.9407/2.7216 - 845.827*eta/6.048 + 415.51*eta2/8.64) + \
          Xisum*(158.0239*eta/5.4432 - 451.597*eta2/6.048 + 20.45*eta2*eta/4.32 + 107.*eta*Xi2/6. - 5.*eta2*Xiprod/24.) + \
          12.*lal.PI*Xi2 - Xi2*Xi*(150.5/2.4 + eta/8.) + \
          Xi*Xiprod*(10.1*eta/2.4 + 3.*eta2/8.)


    AN = 8.*eta*sqrt(lal.PI/5.)
    A2 = (-107. + 55.*eta)/42.
    A3 = 2.*lal.PI - 4.*Xi/3. + 2.*eta*Xisum/3.
    A4 = -2.173/1.512 - eta*(10.69/2.16 - 2.*Xiprod) + 2.047*eta2/1.512
    A5 = -10.7*lal.PI/2.1 + eta*(3.4*lal.PI/2.1)

    A5imag = -24.*eta

    A6 = 270.27409/6.46800 - 8.56*lal.GAMMA/1.05 + \
      2.*lal.PI*lal.PI/3. + \
      eta*(4.1*lal.PI*lal.PI/9.6 - 27.8185/3.3264) - \
      20.261*eta2/2.772 + 11.4635*eta*eta2/9.9792 - \
      4.28*log(16.)/1.05

    A6log = -428./105.

    A6imag = 4.28*lal.PI/1.05

    ### Define other parameters needed by waveform generation ###
    kmin = int(f_min / delta_f)
    kmax = int(f_max / delta_f)
    n = kmax + 1;

    if not out:
        htilde = FrequencySeries(zeros(n,dtype=numpy.complex128), delta_f=delta_f, copy=False)
    else:
        if type(out) is not Array:
            raise TypeError("Output must be an instance of Array")
        if len(out) < kmax:
            raise TypeError("Output array is too small")
        if out.dtype != complex64:
            raise TypeError("Output array is the wrong dtype")
        htilde = FrequencySeries(out, delta_f=delta_f, copy=False)

    phenomC_kernel(htilde.data[kmin:kmax], kmin, delta_f, eta, Xi, distance,
                                       m_sec,  piM,  Mfrd,
                                       pfaN,  pfa2,  pfa3,  pfa4, pfa5,  pfa6,  pfa6log,  pfa7,
                                       a1,  a2,  a3,  a4, a5,  a6,  b1,  b2,
                                       Mf1,  Mf2,  Mf0, d1,  d2,  d0,
                                       xdota2,  xdota3,  xdota4, xdota5,  xdota6,  xdota6log,
                                       xdota7,  xdotaN,  AN, A2,  A3,  A4,  A5,
                                       A5imag,  A6,  A6log,  A6imag,
                                       g1,  del1,  del2,  Q )
    hp = htilde
    hc = htilde * 1j
    return hp, hc
1315532
def g_parameter(time_residual):
    """"""
    mean = np.mean(time_residual)
    time_residual_prime = (time_residual - np.ones(time_residual.shape) * mean)
    time_residual_prime *= time_residual_prime / (-2 * 1.5 * 1.5)
    time_residual_prime = np.exp(time_residual_prime)
    g = np.sum(time_residual_prime) / len(time_residual)
    return g
518588
def _backward_step(self, deltaY, deltaY_sigma2inv_rho_weightT,
                       sigma2_e, weight, mu, mu_Gamma_inv, Gamma_inv,
                       Lambda_0, Lambda_1, H):
        
        n_T = len(Gamma_inv)
        # All the terms with hat before are parameters of posterior
        # distributions of X conditioned on data from all time points,
        # whereas the ones without hat calculated by _forward_step
        # are mean and covariance of posterior of X conditioned on
        # data up to the time point.
        Gamma_inv_hat = [None] * n_T
        mu_Gamma_inv_hat = [None] * n_T
        mu_hat = [None] * n_T

        mu_hat[-1] = mu[-1].copy()
        mu_Gamma_inv_hat[-1] = mu_Gamma_inv[-1].copy()
        Gamma_inv_hat[-1] = Gamma_inv[-1].copy()

        for t in np.arange(n_T - 2, -1, -1):
            tmp = np.linalg.solve(Gamma_inv_hat[t + 1] - Gamma_inv[t + 1]
                                  + Lambda_1, H)
            Gamma_inv_hat[t] = Gamma_inv[t] + Lambda_0 - np.dot(H.T, tmp)
            mu_Gamma_inv_hat[t] = mu_Gamma_inv[t] \
                - deltaY_sigma2inv_rho_weightT[t, :] + np.dot(
                    mu_Gamma_inv_hat[t + 1] - mu_Gamma_inv[t + 1]
                    + np.dot(deltaY[t, :] / sigma2_e, weight.T), tmp)
            mu_hat[t] = np.linalg.solve(Gamma_inv_hat[t],
                                        mu_Gamma_inv_hat[t])
        return mu_hat, mu_Gamma_inv_hat, Gamma_inv_hat
1257925
def iu1transform(self, zeta):
        """"""
        x = real(zeta)
        y = imag(zeta)
        z = x*self.radius.real + y*self.radius.imag
        return self.rot_matrix*z + self.center
17299
def spm(x, model, y=None, n_samples=None, dx_min=-0.1,
        dx_max=0.1, n_dxs=5, dy_min=-0.1, dy_max=0.1, n_dys=5,
        angle_min=-30, angle_max=30, n_angles=31, black_border_size=0):
  
  if y is None:
    preds = model.get_probs(x)
    # Using model predictions as ground truth to avoid label leaking
    preds_max = reduce_max(preds, 1, keepdims=True)
    y = tf.to_float(tf.equal(preds, preds_max))
    y = tf.stop_gradient(y)
    del preds
  y = y / reduce_sum(y, 1, keepdims=True)

  # Define the range of transformations
  dxs = np.linspace(dx_min, dx_max, n_dxs)
  dys = np.linspace(dy_min, dy_max, n_dys)
  angles = np.linspace(angle_min, angle_max, n_angles)

  if n_samples is None:
    import itertools
    transforms = list(itertools.product(*[dxs, dys, angles]))
  else:
    sampled_dxs = np.random.choice(dxs, n_samples)
    sampled_dys = np.random.choice(dys, n_samples)
    sampled_angles = np.random.choice(angles, n_samples)
    transforms = zip(sampled_dxs, sampled_dys, sampled_angles)
  transformed_ims = parallel_apply_transformations(
      x, transforms, black_border_size)

  def _compute_xent(x):
    preds = model.get_logits(x)
    return tf.nn.softmax_cross_entropy_with_logits_v2(
        labels=y, logits=preds)

  all_xents = tf.map_fn(
      _compute_xent,
      transformed_ims,
      parallel_iterations=1)  # Must be 1 to avoid keras race conditions

  # Return the adv_x with worst accuracy

  # all_xents is n_total_samples x batch_size (SB)
  all_xents = tf.stack(all_xents)  # SB

  # We want the worst case sample, with the largest xent_loss
  worst_sample_idx = tf.argmax(all_xents, axis=0)  # B

  batch_size = tf.shape(x)[0]
  keys = tf.stack([
      tf.range(batch_size, dtype=tf.int32),
      tf.cast(worst_sample_idx, tf.int32)
  ], axis=1)
  transformed_ims_bshwc = tf.einsum('sbhwc->bshwc', transformed_ims)
  after_lookup = tf.gather_nd(transformed_ims_bshwc, keys)  # BHWC
  return after_lookup