Combine 1-D rules into multivariate rule using tensor product.
350567
def create_tensorprod_function(funcs):
    """"""
    dim = len(funcs)

    def tensprod_rule(order, part=None):
        """Tensor product rule."""
        order = order*numpy.ones(dim, int)
        values = [funcs[idx](order[idx]) for idx in range(dim)]

        abscissas = [numpy.array(_[0]).flatten() for _ in values]
        abscissas = chaospy.quad.combine(abscissas, part=part).T

        weights = [numpy.array(_[1]).flatten() for _ in values]
        weights = numpy.prod(chaospy.quad.combine(weights, part=part), -1)

        return abscissas, weights

    return tensprod_rule
758850
def fit(self, x, y):
        
        verbose = self.verbose
        # Create temporary files
        data_file = tempfile.NamedTemporaryFile("w+b", delete=False)
        label_file = tempfile.NamedTemporaryFile("w+b", delete=False)

        start = time.time()
        raw_rules = categorical2pysbrl_data(x, y, data_file.name, label_file.name, supp=self.min_support,
                                            zmin=self.min_rule_len, zmax=self.max_rule_len, method=self.fim_method)
        if verbose > 1:
            print("Info: sbrl data files saved to %s and %s temporarily" % (data_file.name, label_file.name))
        data_file.close()
        label_file.close()
        cat_time = time.time() - start
        if verbose:
            print("Info: time for rule mining: %.4fs" % cat_time)
        n_labels = int(np.max(y)) + 1
        start = time.time()
        _model = train_sbrl(data_file.name, label_file.name, self.lambda_, eta=self.eta,
                            max_iters=self.iters, n_chains=self.n_chains, seed=self.seed,
                            alpha=self.alpha, verbose=verbose)
        train_time = time.time() - start
        if verbose:
            print("Info: training time: %.4fs" % train_time)

        # update model parameters
        self._n_classes = n_labels
        self._n_features = x.shape[1]

        # convert the raw parameters to rules
        self.from_raw(_model[0], _model[1], raw_rules)
        self._supports = self.compute_support(x, y)

        # Close the temp files
        os.unlink(data_file.name)
        os.unlink(label_file.name)
1238132
def use_to_ned(tensor):
    '''
    
    '''
    return np.array(ROT_NED_USE.T * np.matrix(tensor) * ROT_NED_USE)
755164
def conv(input, kernel, biases, k_h, k_w, c_o, s_h, s_w,  padding="VALID", group=1):
    
    c_i = input.get_shape()[-1]
    assert c_i%group==0
    assert c_o%group==0
    convolve = lambda i, k: tf.nn.conv2d(i, k, [1, s_h, s_w, 1], padding=padding)
    
    if group==1:
        conv = convolve(input, kernel)
    else:
        input_groups = tf.split(input, group, axis=3)
        kernel_groups = tf.split(kernel, group, axis=3)
        output_groups = [convolve(i, k) for i,k in zip(input_groups, kernel_groups)]
        conv = tf.concat(output_groups, axis=3)
    return  tf.reshape(tf.nn.bias_add(conv, biases), [-1]+conv.get_shape().as_list()[1:])
796001
def toy_model(features, mesh):
  """"""
  batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
  io_dim = mtf.Dimension('io', FLAGS.io_size)

  master_dtype = tf.as_dtype(FLAGS.master_dtype)
  slice_dtype = tf.as_dtype(FLAGS.slice_dtype)
  activation_dtype = tf.as_dtype(FLAGS.activation_dtype)

  x = mtf.import_tf_tensor(mesh, features, mtf.Shape([batch_dim, io_dim]))
  x = mtf.cast(x, activation_dtype)
  h = x
  for lnum in xrange(1, FLAGS.num_hidden_layers + 2):
    if lnum + 1 == FLAGS.num_hidden_layers + 2:
      # output layer
      dim = io_dim
    elif lnum % 2 == 0:
      dim = mtf.Dimension('hidden_even', FLAGS.hidden_size)
    else:
      dim = mtf.Dimension('hidden_odd', FLAGS.hidden_size)
    h = mtf.layers.dense(
        h, dim,
        use_bias=False,
        master_dtype=master_dtype,
        slice_dtype=slice_dtype,
        name='layer_%d' % lnum)
  y = h
  g = tf.train.get_global_step()
  if FLAGS.step_with_nan >= 0:
    # Trigger NaN in the forward pass, this is used for testing whether
    # MeshTensorFlow can handle occasional NaN value.
    y += mtf.import_tf_tensor(
        mesh,
        tf.divide(
            0.0,
            tf.cond(tf.equal(g, FLAGS.step_with_nan), lambda: 0., lambda: 1.)),
        mtf.Shape([]))

  loss = mtf.reduce_mean(mtf.square(y - x))
  return y, loss
258198
def _matmul_with_relative_keys_2d(x, y, heads_share_relative_embedding):
  """"""
  if heads_share_relative_embedding:
    ret = tf.einsum("bhxyd,md->bhxym", x, y)
  else:
    ret = tf.einsum("bhxyd,hmd->bhxym", x, y)
  return ret
879443
def conv1d_block(inputs, filters, dilation_rates_and_kernel_sizes, **kwargs):
  """"""
  return conv_block_internal(conv1d, inputs, filters,
                             dilation_rates_and_kernel_sizes, **kwargs)
126185
def make_prediction_output_tensors(args, features, input_ops, model_fn_ops,
                                   keep_target):
  """"""
  target_name = feature_transforms.get_target_name(features)
  key_names = get_key_names(features)

  outputs = {}
  outputs.update({key_name: tf.squeeze(input_ops.features[key_name])
                  for key_name in key_names})

  if is_classification_model(args.model):

    # build maps from ints to the origional categorical strings.
    class_names = read_vocab(args, target_name)
    table = tf.contrib.lookup.index_to_string_table_from_tensor(
        mapping=class_names,
        default_value='UNKNOWN')

    # Get the label of the input target.
    if keep_target:
      input_target_label = table.lookup(input_ops.features[target_name])
      outputs[PG_TARGET] = tf.squeeze(input_target_label)

    # TODO(brandondutra): get the score of the target label too.
    probabilities = model_fn_ops.predictions['probabilities']

    # if top_n == 0, this means use all the classes. We will use class names as
    # probabilities labels.
    if args.top_n == 0:
      predicted_index = tf.argmax(probabilities, axis=1)
      predicted = table.lookup(predicted_index)
      outputs.update({PG_CLASSIFICATION_FIRST_LABEL: predicted})
      probabilities_list = tf.unstack(probabilities, axis=1)
      for class_name, p in zip(class_names, probabilities_list):
        outputs[class_name] = p
    else:
      top_n = args.top_n

      # get top k labels and their scores.
      (top_k_values, top_k_indices) = tf.nn.top_k(probabilities, k=top_n)
      top_k_labels = table.lookup(tf.to_int64(top_k_indices))

      # Write the top_k values using 2*top_n columns.
      num_digits = int(math.ceil(math.log(top_n, 10)))
      if num_digits == 0:
        num_digits = 1
      for i in range(0, top_n):
        # Pad i based on the size of k. So if k = 100, i = 23 -> i = '023'. This
        # makes sorting the columns easy.
        padded_i = str(i + 1).zfill(num_digits)

        if i == 0:
          label_alias = PG_CLASSIFICATION_FIRST_LABEL
        else:
          label_alias = PG_CLASSIFICATION_LABEL_TEMPLATE % padded_i

        label_tensor_name = (tf.squeeze(
            tf.slice(top_k_labels, [0, i], [tf.shape(top_k_labels)[0], 1])))

        if i == 0:
          score_alias = PG_CLASSIFICATION_FIRST_SCORE
        else:
          score_alias = PG_CLASSIFICATION_SCORE_TEMPLATE % padded_i

        score_tensor_name = (tf.squeeze(
            tf.slice(top_k_values,
                     [0, i],
                     [tf.shape(top_k_values)[0], 1])))

        outputs.update({label_alias: label_tensor_name,
                        score_alias: score_tensor_name})

  else:
    if keep_target:
      outputs[PG_TARGET] = tf.squeeze(input_ops.features[target_name])

    scores = model_fn_ops.predictions['scores']
    outputs[PG_REGRESSION_PREDICTED_TARGET] = tf.squeeze(scores)

  return outputs
1289088
def deconv2d(
    input_, output_shape, k_h, k_w, d_h, d_w, stddev=0.02, name="deconv2d"):
  """"""
  with tf.variable_scope(name):
    w = tf.get_variable(
        "w", [k_h, k_w, output_shape[-1], input_.get_shape()[-1]],
        initializer=tf.random_normal_initializer(stddev=stddev))
    deconv = tf.nn.conv2d_transpose(
        input_, w, output_shape=output_shape, strides=[1, d_h, d_w, 1])
    biases = tf.get_variable(
        "biases", [output_shape[-1]], initializer=tf.constant_initializer(0.0))
    return tf.reshape(tf.nn.bias_add(deconv, biases), deconv.get_shape())
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
54959
def susvd(x, x_obs, rho, penalties):
    

    assert type(x) == Tensor, "Input array must be a Tensor"

    while True:

        # proximal operator for the Fro. norm
        x = squared_error(x, rho, x_obs)

        # sequential singular value thresholding
        for ix, penalty in enumerate(penalties):
            x = x.unfold(ix).svt(penalty / rho).fold()

        yield x
449152
def multifuncrs(X, funs, eps=1E-6,
                nswp=10,
                kickrank=5,
                y0=None,
                rmax=999999,  # TODO:infinity \
                kicktype='amr-two',        \
                pcatype='svd',             \
                trunctype='fro',           \
                d2=1,                      \
                do_qr=False,               \
                verb=1):
    

    dtype = np.float64
    if len([x for x in X if x.is_complex]) > 0:
        dtype = np.complex128

    y = y0
    wasrand = False

    nx = len(X)
    d = X[0].d
    n = X[0].n
    rx = np.transpose(np.array([ttx.r for ttx in X]))
    #crx = [tt.tensor.to_list(ttx) for x in X]
    #crx = zip(*crx)
    crx = np.transpose(np.array([tt.tensor.to_list(ttx)
                                 for ttx in X], dtype=np.object))
    crx = np.empty((nx, d), dtype=np.object)
    i = 0
    for ttx in X:
        v = tt.tensor.to_list(ttx)
        j = 0
        for w in v:
            crx[i, j] = w
            j = j + 1
        i = i + 1
    crx = crx.T
    if y is None:
        ry = d2 * np.ones((d + 1,), dtype=np.int32)
        ry[0] = 1
        y = tt.rand(n, d, ry)
        wasrand = True

    ry = y.r
    cry = tt.tensor.to_list(y)

    Ry = np.zeros((d + 1, ), dtype=np.object)
    Ry[0] = np.array([[1.0]], dtype=dtype)
    Ry[d] = np.array([[1.0]], dtype=dtype)
    Rx = np.zeros((d + 1, nx), dtype=np.object)
    Rx[0, :] = np.ones(nx, dtype=dtype)
    Rx[d, :] = np.ones(nx, dtype=dtype)

    block_order = [+d, -d]

    # orth
    for i in range(0, d - 1):
        cr = cry[i]
        cr = reshape(cr, (ry[i] * n[i], ry[i + 1]))
        cr, rv = np.linalg.qr(cr)
        cr2 = cry[i + 1]
        cr2 = reshape(cr2, (ry[i + 1], n[i + 1] * ry[i + 2]))
        cr2 = np.dot(rv, cr2)  # matrix multiplication
        ry[i + 1] = cr.shape[1]
        cr = reshape(cr, (ry[i], n[i], ry[i + 1]))
        cry[i + 1] = reshape(cr2, (ry[i + 1], n[i + 1], ry[i + 2]))
        cry[i] = cr

        Ry[i + 1] = np.dot(Ry[i], reshape(cr, (ry[i], n[i] * ry[i + 1])))
        Ry[i + 1] = reshape(Ry[i + 1], (ry[i] * n[i], ry[i + 1]))
        curind = []
        if wasrand:
            # EVERY DAY I'M SHUFFLIN'
            curind = np.random.permutation(n[i] * ry[i])[:ry[i + 1]]
        else:
            curind = maxvol(Ry[i + 1])
        Ry[i + 1] = Ry[i + 1][curind, :]
        for j in range(0, nx):
            try:
                Rx[i + 1, j] = reshape(crx[i, j],
                                       (rx[i, j], n[i] * rx[i + 1, j]))
            except:
                pass
            Rx[i + 1, j] = np.dot(Rx[i, j], Rx[i + 1, j])
            Rx[i + 1, j] = reshape(Rx[i + 1, j], (ry[i] * n[i], rx[i + 1, j]))
            Rx[i + 1, j] = Rx[i + 1, j][curind, :]

    d2 = ry[d]
    ry[d] = 1
    cry[d - 1] = np.transpose(cry[d - 1], [2, 0, 1])  # permute

    last_sweep = False
    swp = 1

    dy = np.zeros((d, ))
    max_dy = 0

    cur_order = copy.copy(block_order)
    order_index = 1
    i = d - 1
    # can't use 'dir' identifier in python
    dirn = int(math.copysign(1, cur_order[order_index]))

    # DMRG sweeps
    while swp <= nswp or dirn > 0:

        oldy = reshape(cry[i], (d2 * ry[i] * n[i] * ry[i + 1],))

        if not last_sweep:
            # compute the X superblocks
            curbl = np.zeros((ry[i] * n[i] * ry[i + 1], nx), dtype=dtype)
            for j in range(0, nx):
                cr = reshape(crx[i, j], (rx[i, j], n[i] * rx[i + 1, j]))
                cr = np.dot(Rx[i, j], cr)
                cr = reshape(cr, (ry[i] * n[i], rx[i + 1, j]))
                cr = np.dot(cr, Rx[i + 1, j])
                curbl[:, j] = cr.flatten('F')
            # call the function
            newy = funs(curbl)
            # multiply with inverted Ry
            newy = reshape(newy, (ry[i], n[i] * ry[i + 1] * d2))
            newy = np.linalg.solve(Ry[i], newy)  # y = R \ y
            newy = reshape(newy, (ry[i] * n[i] * ry[i + 1], d2))
            newy = reshape(np.transpose(newy), (d2 * ry[i] * n[i], ry[i + 1]))
            newy = np.transpose(np.linalg.solve(
                np.transpose(Ry[i + 1]), np.transpose(newy)))  # y=y/R
            newy = reshape(newy, (d2 * ry[i] * n[i] * ry[i + 1],))
        else:
            newy = oldy

        dy[i] = np.linalg.norm(newy - oldy) / np.linalg.norm(newy)
        max_dy = max(max_dy, dy[i])

        # truncation
        if dirn > 0:  # left-to-right
            newy = reshape(newy, (d2, ry[i] * n[i] * ry[i + 1]))
            newy = reshape(np.transpose(newy), (ry[i] * n[i], ry[i + 1] * d2))
        else:
            newy = reshape(newy, (d2 * ry[i], n[i] * ry[i + 1]))

        r = 0  # defines a variable in global scope

        if kickrank >= 0:
            u, s, v = np.linalg.svd(newy, full_matrices=False)
            v = np.conj(np.transpose(v))
            if trunctype == "fro" or last_sweep:
                r = my_chop2(s, eps / math.sqrt(d) * np.linalg.norm(s))
            else:
                # truncate taking into account the (r+1) overhead in the cross
                # (T.S.: what?)
                cums = abs(s * np.arange(2, len(s) + 2)) ** 2
                cums = np.cumsum(cums[::-1])[::-1]
                cums = cums / cums[0]
                ff = [i for i in range(len(cums)) if cums[i] < eps ** 2 / d]
                if len(ff) == 0:
                    r = len(s)
                else:
                    r = np.amin(ff)
            r = min(r, rmax, len(s))
        else:
            if dirn > 0:
                u, v = np.linalg.qr(newy)
                v = np.conj(np.transpose(v))
                r = u.shape[1]
                s = np.ones((r, ))
            else:
                v, u = np.linalg.qr(np.transpose(newy))
                v = np.conj(v)
                u = np.transpose(u)
                r = u.shape[1]
                s = np.ones((r, ))

        if verb > 1:
            print('=multifuncrs=   block %d{%d}, dy: %3.3e, r: %d' % (i, dirn, dy[i], r))

        # kicks and interfaces
        if dirn > 0 and i < d - 1:
            u = u[:, :r]
            v = np.dot(v[:, :r], np.diag(s[:r]))

            # kick
            radd = 0
            rv = 1
            if not last_sweep and kickrank > 0:
                uk = None
                if kicktype == 'amr-two':
                    # AMR(two)-like kick.

                    # compute the X superblocks
                    ind2 = np.unique(np.random.randint(
                        0, ry[i + 2] * n[i + 1], ry[i + 1]))
                    #ind2 = np.unique(np.floor(np.random.rand(ry[i + 1]) * (ry[i + 2] * n[i + 1])))
                    rkick = len(ind2)
                    curbl = np.zeros((ry[i] * n[i] * rkick, nx), dtype=dtype)
                    for j in range(nx):
                        cr1 = reshape(
                            crx[i, j], (rx[i, j], n[i] * rx[i + 1, j]))
                        cr1 = np.dot(Rx[i, j], cr1)
                        cr1 = reshape(cr1, (ry[i] * n[i], rx[i + 1, j]))
                        cr2 = reshape(
                            crx[i + 1, j], (rx[i + 1, j] * n[i + 1], rx[i + 2, j]))
                        cr2 = np.dot(cr2, Rx[i + 2, j])
                        cr2 = reshape(
                            cr2, (rx[i + 1, j], n[i + 1] * ry[i + 2]))
                        cr2 = cr2[:, ind2]
                        curbl[:, j] = reshape(
                            np.dot(cr1, cr2), (ry[i] * n[i] * rkick,))
                    # call the function
                    uk = funs(curbl)
                    uk = reshape(uk, (ry[i], n[i] * rkick * d2))
                    uk = np.linalg.solve(Ry[i], uk)
                    uk = reshape(uk, (ry[i] * n[i], rkick * d2))
                    if pcatype == 'svd':
                        uk, sk, vk = np.linalg.svd(uk, full_matrices=False)
                        vk = np.conj(np.transpose(vk))
                        uk = uk[:, :min(kickrank, uk.shape[1])]
                    else:
                        # uk = uchol(np.transpose(uk), kickrank + 1) # TODO
                        uk = uk[:, :max(uk.shape[1] - kickrank + 1, 1):-1]
                else:
                    uk = np.random.rand(ry[i] * n[i], kickrank)
                u, rv = np.linalg.qr(np.concatenate((u, uk), axis=1))
                radd = uk.shape[1]
            v = np.concatenate(
                (v, np.zeros((ry[i + 1] * d2, radd), dtype=dtype)), axis=1)
            v = np.dot(rv, np.conj(np.transpose(v)))
            r = u.shape[1]

            cr2 = cry[i + 1]
            cr2 = reshape(cr2, (ry[i + 1], n[i + 1] * ry[i + 2]))
            v = reshape(v, (r * ry[i + 1], d2))
            v = reshape(np.transpose(v), (d2 * r, ry[i + 1]))
            v = np.dot(v, cr2)

            ry[i + 1] = r

            u = reshape(u, (ry[i], n[i], r))
            v = reshape(v, (d2, r, n[i + 1], ry[i + 2]))

            cry[i] = u
            cry[i + 1] = v

            Ry[i + 1] = np.dot(Ry[i], reshape(u, (ry[i], n[i] * ry[i + 1])))
            Ry[i + 1] = reshape(Ry[i + 1], (ry[i] * n[i], ry[i + 1]))
            curind = maxvol(Ry[i + 1])
            Ry[i + 1] = Ry[i + 1][curind, :]
            for j in range(nx):
                Rx[i + 1, j] = reshape(crx[i, j],
                                       (rx[i, j], n[i] * rx[i + 1, j]))
                Rx[i + 1, j] = np.dot(Rx[i, j], Rx[i + 1, j])
                Rx[i + 1, j] = reshape(Rx[i + 1, j],
                                       (ry[i] * n[i], rx[i + 1, j]))
                Rx[i + 1, j] = Rx[i + 1, j][curind, :]
        elif dirn < 0 and i > 0:
            u = np.dot(u[:, :r], np.diag(s[:r]))
            v = np.conj(v[:, :r])

            radd = 0
            rv = 1
            if not last_sweep and kickrank > 0:
                if kicktype == 'amr-two':
                    # compute the X superblocks
                    ind2 = np.unique(np.random.randint(
                        0, ry[i - 1] * n[i - 1], ry[i]))
                    rkick = len(ind2)
                    curbl = np.zeros(
                        (rkick * n[i] * ry[i + 1], nx), dtype=dtype)
                    for j in range(nx):
                        cr1 = reshape(
                            crx[i, j], (rx[i, j] * n[i], rx[i + 1, j]))
                        cr1 = np.dot(cr1, Rx[i + 1, j])
                        cr1 = reshape(cr1, (rx[i, j], n[i] * ry[i + 1]))
                        cr2 = reshape(
                            crx[i - 1, j], (rx[i - 1, j], n[i - 1] * rx[i, j]))
                        cr2 = np.dot(Rx[i - 1, j], cr2)
                        cr2 = reshape(cr2, (ry[i - 1] * n[i - 1], rx[i, j]))
                        cr2 = cr2[ind2, :]
                        curbl[:, j] = reshape(
                            np.dot(cr2, cr1), (rkick * n[i] * ry[i + 1],))
                    # calling the function
                    uk = funs(curbl)
                    uk = reshape(uk, (rkick * n[i] * ry[i + 1], d2))
                    uk = reshape(np.transpose(
                        uk), (d2 * rkick * n[i], ry[i + 1]))
                    uk = np.transpose(np.linalg.solve(
                        np.transpose(Ry[i + 1]), np.transpose(uk)))
                    uk = reshape(uk, (d2 * rkick, n[i] * ry[i + 1]))
                    if pcatype == 'svd':
                        vk, sk, uk = np.linalg.svd(uk, full_matrices=False)
                        uk = np.conj(np.transpose(uk))
                        # TODO: refactor
                        uk = uk[:, :min(kickrank, uk.shape[1])]
                    else:
                        # uk = uchol(uk, kickrank + 1) # TODO
                        uk = uk[:, :max(uk.shape[1] - kickrank + 1, 1):-1]
                else:
                    uk = np.random.rand(n[i] * ry[i + 1], kickrank)
                v, rv = np.linalg.qr(np.concatenate((v, uk), axis=1))
                radd = uk.shape[1]
            u = np.concatenate(
                (u, np.zeros((d2 * ry[i], radd), dtype=dtype)), axis=1)
            u = np.dot(u, np.transpose(rv))
            r = v.shape[1]
            cr2 = cry[i - 1]
            cr2 = reshape(cr2, (ry[i - 1] * n[i - 1], ry[i]))
            u = reshape(u, (d2, ry[i] * r))
            u = reshape(np.transpose(u), (ry[i], r * d2))
            u = np.dot(cr2, u)

            u = reshape(u, (ry[i - 1] * n[i - 1] * r, d2))
            u = reshape(np.transpose(u), (d2, ry[i - 1], n[i - 1], r))
            v = reshape(np.transpose(v), (r, n[i], ry[i + 1]))

            ry[i] = r
            cry[i - 1] = u
            cry[i] = v

            Ry[i] = np.dot(reshape(v, (ry[i] * n[i], ry[i + 1])), Ry[i + 1])
            Ry[i] = reshape(Ry[i], (ry[i], n[i] * ry[i + 1]))
            curind = maxvol(np.transpose(Ry[i]))
            Ry[i] = Ry[i][:, curind]
            for j in range(nx):
                Rx[i, j] = reshape(crx[i, j], (rx[i, j] * n[i], rx[i + 1, j]))
                Rx[i, j] = np.dot(Rx[i, j], Rx[i + 1, j])
                Rx[i, j] = reshape(Rx[i, j], (rx[i, j], n[i] * ry[i + 1]))
                Rx[i, j] = Rx[i, j][:, curind]
        elif dirn > 0 and i == d - 1:
            newy = np.dot(np.dot(u[:, :r], np.diag(s[:r])),
                          np.conj(np.transpose(v[:, :r])))
            newy = reshape(newy, (ry[i] * n[i] * ry[i + 1], d2))
            cry[i] = reshape(np.transpose(newy), (d2, ry[i], n[i], ry[i + 1]))
        elif dirn < 0 and i == 0:
            newy = np.dot(np.dot(u[:, :r], np.diag(s[:r])),
                          np.conj(np.transpose(v[:, :r])))
            newy = reshape(newy, (d2, ry[i], n[i], ry[i + 1]))
            cry[i] = newy

        i = i + dirn
        cur_order[order_index] = cur_order[order_index] - dirn
        if cur_order[order_index] == 0:
            order_index = order_index + 1
            if verb > 0:
                print('=multifuncrs= sweep %d{%d}, max_dy: %3.3e, erank: %g' % (swp, order_index, max_dy,
                                                                                math.sqrt(np.dot(ry[:d], n * ry[1:]) / np.sum(n))))

            if last_sweep:
                break
            if max_dy < eps and dirn < 0:
                last_sweep = True
                kickrank = 0

            if order_index >= len(cur_order):
                cur_order = copy.copy(block_order)
                order_index = 0
                if last_sweep:
                    cur_order = [d - 1]

                max_dy = 0
                swp = swp + 1

            dirn = int(math.copysign(1, cur_order[order_index]))
            i = i + dirn

    cry[d - 1] = np.transpose(cry[d - 1][:, :, :, 0], [1, 2, 0])
    y = tt.tensor.from_list(cry)
    return y
448535
def combine_initial_dims(tensor: torch.Tensor) -> torch.Tensor:
    
    if tensor.dim() <= 2:
        return tensor
    else:
        return tensor.view(-1, tensor.size(-1))
1034900
def model_fn(features, labels, mode, params):
  """"""
  tf.logging.info("features = %s labels = %s mode = %s params=%s" %
                  (features, labels, mode, params))
  global_step = tf.train.get_global_step()
  graph = mtf.Graph()
  mesh = mtf.Mesh(graph, "my_mesh")
  logits, loss = mnist_model(features, labels, mesh)
  mesh_shape = mtf.convert_to_shape(FLAGS.mesh_shape)
  layout_rules = mtf.convert_to_layout_rules(FLAGS.layout)
  mesh_size = mesh_shape.size
  mesh_devices = [""] * mesh_size
  mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
      mesh_shape, layout_rules, mesh_devices)

  if mode == tf.estimator.ModeKeys.TRAIN:
    var_grads = mtf.gradients(
        [loss], [v.outputs[0] for v in graph.trainable_variables])
    optimizer = mtf.optimize.AdafactorOptimizer()
    update_ops = optimizer.apply_grads(var_grads, graph.trainable_variables)

  lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  restore_hook = mtf.MtfRestoreHook(lowering)

  tf_logits = lowering.export_to_tf_tensor(logits)
  if mode != tf.estimator.ModeKeys.PREDICT:
    tf_loss = lowering.export_to_tf_tensor(loss)
    tf.summary.scalar("loss", tf_loss)

  if mode == tf.estimator.ModeKeys.TRAIN:
    tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
    tf_update_ops.append(tf.assign_add(global_step, 1))
    train_op = tf.group(tf_update_ops)
    saver = tf.train.Saver(
        tf.global_variables(),
        sharded=True,
        max_to_keep=10,
        keep_checkpoint_every_n_hours=2,
        defer_build=False, save_relative_paths=True)
    tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
    saver_listener = mtf.MtfCheckpointSaverListener(lowering)
    saver_hook = tf.train.CheckpointSaverHook(
        FLAGS.model_dir,
        save_steps=1000,
        saver=saver,
        listeners=[saver_listener])

    accuracy = tf.metrics.accuracy(
        labels=labels, predictions=tf.argmax(tf_logits, axis=1))

    # Name tensors to be logged with LoggingTensorHook.
    tf.identity(tf_loss, "cross_entropy")
    tf.identity(accuracy[1], name="train_accuracy")

    # Save accuracy scalar to Tensorboard output.
    tf.summary.scalar("train_accuracy", accuracy[1])

    # restore_hook must come before saver_hook
    return tf.estimator.EstimatorSpec(
        tf.estimator.ModeKeys.TRAIN, loss=tf_loss, train_op=train_op,
        training_chief_hooks=[restore_hook, saver_hook])

  if mode == tf.estimator.ModeKeys.PREDICT:
    predictions = {
        "classes": tf.argmax(tf_logits, axis=1),
        "probabilities": tf.nn.softmax(tf_logits),
    }
    return tf.estimator.EstimatorSpec(
        mode=tf.estimator.ModeKeys.PREDICT,
        predictions=predictions,
        prediction_hooks=[restore_hook],
        export_outputs={
            "classify": tf.estimator.export.PredictOutput(predictions)
        })
  if mode == tf.estimator.ModeKeys.EVAL:
    return tf.estimator.EstimatorSpec(
        mode=tf.estimator.ModeKeys.EVAL,
        loss=tf_loss,
        evaluation_hooks=[restore_hook],
        eval_metric_ops={
            "accuracy":
            tf.metrics.accuracy(
                labels=labels, predictions=tf.argmax(tf_logits, axis=1)),
        })
241802
def tensor_product_basis(*bases):
    
    dim = np.prod([basis.data.shape[1] for basis in bases])
    tp_basis = np.zeros((dim**2, dim, dim), dtype=complex)

    for idx_factors, factors in enumerate(it.product(*[basis.data for basis in bases])):
        tp_basis[idx_factors, :, :] = reduce(np.kron, factors)

    return TomographyBasis(tp_basis,
        sum((
            factor.dims for factor in bases
        ), []),
        list(map(
        r"\otimes".join,
        it.product(*[
            basis.labels for basis in bases
        ])
    )))
203757
def multiply(x1, x2, output_shape=None, name=None):
  
  if not isinstance(x2, Tensor):
    return ScalarMultiplyOperation(x1, x2).outputs[0]
  with tf.name_scope(name, default_name="mul"):
    x1, x2 = binary_arguments_to_tensors(x1, x2)
    return einsum(
        [x1, x2],
        output_shape=_infer_binary_broadcast_shape(
            x1.shape, x2.shape, output_shape))
774074
def _divide(self, x1, x2, out):
        """"""
        self.tspace._divide(x1.tensor, x2.tensor, out.tensor)