Helper function for local 2d attention.

  Args:
    x: a [batch, height, width, depth] tensor
    block_h: An integer. block height
    block_w: An inteter. block width

  returns:
    a [batch, num_heads, height/block_h, width/block_w, depth] tensor
483379
def _extract_blocks(x, block_h, block_w):
  
  (_, height, width, depth) = common_layers.shape_list(x)
  assert height % block_h == 0
  assert width % block_w == 0
  x = tf.reshape(x, [-1, height//block_h, block_h,
                     width//block_w, block_w, depth])
  return tf.transpose(x, [0, 1, 3, 2, 4, 5])
238316
def attend(x, source, hparams, name):
  """"""
  with tf.variable_scope(name):
    x = tf.squeeze(x, axis=2)
    if len(source.get_shape()) > 3:
      source = tf.squeeze(source, axis=2)
    source = common_attention.add_timing_signal_1d(source)
    y = common_attention.multihead_attention(
        common_layers.layer_preprocess(x, hparams), source, None,
        hparams.attention_key_channels or hparams.hidden_size,
        hparams.attention_value_channels or hparams.hidden_size,
        hparams.hidden_size, hparams.num_heads,
        hparams.attention_dropout)
    res = common_layers.layer_postprocess(x, y, hparams)
    return tf.expand_dims(res, axis=2)
8119
def chunked_X(self, chunk_size: Optional[int] = None):
        
        if chunk_size is None:
            # Should be some adaptive code
            chunk_size = 6000
        start = 0
        n = self.n_obs
        for _ in range(int(n // chunk_size)):
            end = start + chunk_size
            yield (self.X[start:end], start, end)
            start = end
        if start < n:
            yield (self.X[start:n], start, n)
662653
def linear(self, x):
    
    with tf.name_scope("presoftmax_linear"):
      batch_size = tf.shape(x)[0]
      length = tf.shape(x)[1]

      x = tf.reshape(x, [-1, self.hidden_size])
      logits = tf.matmul(x, self.shared_weights, transpose_b=True)

      return tf.reshape(logits, [batch_size, length, self.vocab_size])
910063
def forward_tensor(self, x):
        
        fwd = vec_to_tri(x, self.N)
        return tf.squeeze(fwd, axis=0) if self.squeeze else fwd
126720
def scatter_blocks_2d(x, indices, shape):
  """"""
  x_shape = common_layers.shape_list(x)
  # [length, batch, heads, dim]
  x_t = tf.transpose(
      tf.reshape(x, [x_shape[0], x_shape[1], -1, x_shape[-1]]), [2, 0, 1, 3])
  x_t_shape = common_layers.shape_list(x_t)
  indices = tf.reshape(indices, [-1, 1])
  scattered_x = tf.scatter_nd(indices, x_t, x_t_shape)
  scattered_x = tf.transpose(scattered_x, [1, 2, 0, 3])
  return tf.reshape(scattered_x, shape)
912069
def _make_local_block(x, depth, batch, heads, num_blocks, block_length):
  """"""
  prev_block = tf.slice(x, [0, 0, 0, 0, 0],
                        [-1, -1, num_blocks - 1, -1, -1])
  cur_block = tf.slice(x, [0, 0, 1, 0, 0], [-1, -1, -1, -1, -1])
  local_block = tf.concat([prev_block, cur_block], 3)
  return tf.reshape(local_block,
                    [batch, heads, num_blocks - 1, block_length * 2, depth])
45701
def pad_to_multiple_2d(x, block_shape):
  
  old_shape = x.get_shape().dims
  last = old_shape[-1]
  if len(old_shape) == 4:
    height_padding = -common_layers.shape_list(x)[1] % block_shape[0]
    width_padding = -common_layers.shape_list(x)[2] % block_shape[1]
    paddings = [[0, 0], [0, height_padding], [0, width_padding], [0, 0]]
  elif len(old_shape) == 5:
    height_padding = -common_layers.shape_list(x)[2] % block_shape[0]
    width_padding = -common_layers.shape_list(x)[3] % block_shape[1]
    paddings = [[0, 0], [0, 0], [0, height_padding], [0, width_padding], [0, 0]]

  padded_x = tf.pad(x, paddings)
  padded_shape = padded_x.get_shape().as_list()
  padded_shape = padded_shape[:-1] + [last]
  padded_x.set_shape(padded_shape)
  return padded_x
1075214
def wc_dict2lha(wc, skip_redundant=True, skip_zero=True):
    """"""
    d = OrderedDict()
    for name, (block, i) in WC_dict_0f.items():
        if block not in d:
            d[block] = defaultdict(list)
        if wc[name] != 0:
            d[block]['values'].append([i, wc[name].real])
    for name in definitions.WC_keys_2f:
        reblock = 'WC'+name.upper()
        imblock = 'IMWC'+name.upper()
        if reblock not in d:
            d[reblock] = defaultdict(list)
        if imblock not in d:
            d[imblock] = defaultdict(list)
        for i in range(3):
            for j in range(3):
                if (i, j) in definitions.redundant_elements[name] and skip_redundant:
                    # skip redundant elements
                    continue
                if wc[name][i, j].real != 0 or not skip_zero:
                    d[reblock]['values'].append([i+1, j+1, float(wc[name][i, j].real)])
                if wc[name][i, j].imag != 0 or not skip_zero:
                    # omit Im parts that have to vanish by symmetry
                    if (i, j) not in definitions.vanishing_im_parts[name]:
                        d[imblock]['values'].append([i+1, j+1, float(wc[name][i, j].imag)])
    for name in definitions.WC_keys_4f:
        reblock = 'WC'+name.upper()
        imblock = 'IMWC'+name.upper()
        if reblock not in d:
            d[reblock] = defaultdict(list)
        if imblock not in d:
            d[imblock] = defaultdict(list)
        for i in range(3):
            for j in range(3):
                for k in range(3):
                    for l in range(3):
                        if (i, j, k, l) in definitions.redundant_elements[name] and skip_redundant:
                            # skip redundant elements
                            continue
                        if wc[name][i, j, k, l].real != 0 or not skip_zero:
                            d[reblock]['values'].append([i+1, j+1, k+1, l+1, float(wc[name][i, j, k, l].real)])
                        if wc[name][i, j, k, l].imag != 0 or not skip_zero:
                            # omit Im parts that have to vanish by symmetry
                            if (i, j, k, l) not in definitions.vanishing_im_parts[name]:
                                d[imblock]['values'].append([i+1, j+1, k+1, l+1, float(wc[name][i, j, k, l].imag)])
    # remove empty blocks
    empty = []
    for block in d:
        if d[block] == {}:
            empty.append(block)
    for block in empty:
        del d[block]
    return {'BLOCK': d}
559680
def dilated_attention_1d(x,
                         hparams,
                         attention_type="masked_dilated_1d",
                         q_padding="VALID",
                         kv_padding="VALID",
                         gap_size=2):
  """"""
  # self-attention
  x, x_shape, is_4d = maybe_reshape_4d_to_3d(x)
  with tf.variable_scope("masked_dilated_1d"):
    y = common_attention.multihead_attention(
        x,
        None,
        None,
        hparams.attention_key_channels or hparams.hidden_size,
        hparams.attention_value_channels or hparams.hidden_size,
        hparams.hidden_size,
        hparams.num_heads,
        hparams.attention_dropout,
        attention_type=attention_type,
        block_width=hparams.block_width,
        block_length=hparams.block_length,
        q_padding=q_padding,
        kv_padding=kv_padding,
        q_filter_width=hparams.q_filter_width,
        kv_filter_width=hparams.kv_filter_width,
        gap_size=gap_size,
        num_memory_blocks=hparams.num_memory_blocks,
        name="self_attention")
    if is_4d:
      y = tf.reshape(y, x_shape)
      y.set_shape([None, None, None, hparams.hidden_size])
    return y
632351
public static void computeV_blockVector( final int blockLength ,
                                             final DSubmatrixD1 A ,
                                             final double gammas[] ,
                                             final DSubmatrixD1 V )
    {
        int blockHeight = Math.min(blockLength,A.row1-A.row0);
        if( blockHeight <= 1 )
            return;
        int width = A.col1-A.col0;
        int num = Math.min(width-1,blockHeight);

        for( int i = 0; i < num; i++ ) {
            double gamma = gammas[A.row0+i];

            // compute y
            computeY(blockLength,A,V,i,gamma);

            // compute v from y
            computeRowOfV(blockLength,A,V,i,gamma);
        }
    }
1009217
def compress_self_attention_layer(x, hparams, name=None):
  """"""
  with tf.variable_scope(name, default_name="compress_self_attention"):
    x, xshape, _ = cia.maybe_reshape_4d_to_3d(x)
    y = common_attention.multihead_attention(
        common_layers.layer_preprocess(x, hparams),
        None,
        None,
        hparams.attention_key_channels or hparams.hidden_size,
        hparams.attention_value_channels or hparams.hidden_size,
        hparams.hidden_size, hparams.num_heads,
        hparams.attention_dropout)
    res = common_layers.layer_postprocess(x, y, hparams)
    return tf.reshape(res, xshape)
454292
def get_self_attention_bias(x):
  

  x_shape = common_layers.shape_list(x)
  self_attention_bias = common_attention.attention_bias_lower_triangle(
      x_shape[1])
  return self_attention_bias
878054
def xmoe_tr_2d():
  
  hparams = xmoe_tr_dense_2k()
  hparams.mesh_shape = "b0:2;b1:4"
  hparams.outer_batch_size = 4
  hparams.layout = "outer_batch:b0;inner_batch:b1,expert_x:b1,expert_y:b0"
  hparams.encoder_layers = ["self_att", "moe_2d"] * 4
  hparams.decoder_layers = ["self_att", "enc_att", "moe_2d"] * 4
  hparams.moe_hidden_size = 2048
  hparams.moe_experts_x = 4
  hparams.moe_experts_y = 4
  return hparams
149896
def __realized_bbox(self, requested_bbox):
    
    realized_bbox = requested_bbox.expand_to_chunk_size(self.underlying, offset=self.voxel_offset)
    return Bbox.clamp(realized_bbox, self.bounds)
1365574
def tensor_vstack(tensor_list, pad=0):
    
    if len(tensor_list) == 1:
        return tensor_list[0][np.newaxis, :]

    ndim = len(tensor_list[0].shape)
    dimensions = [len(tensor_list)]  # first dim is batch size
    for dim in range(ndim):
        dimensions.append(max([tensor.shape[dim] for tensor in tensor_list]))

    dtype = tensor_list[0].dtype
    if pad == 0:
        all_tensor = np.zeros(tuple(dimensions), dtype=dtype)
    elif pad == 1:
        all_tensor = np.ones(tuple(dimensions), dtype=dtype)
    else:
        all_tensor = np.full(tuple(dimensions), pad, dtype=dtype)
    if ndim == 1:
        for ind, tensor in enumerate(tensor_list):
            all_tensor[ind, :tensor.shape[0]] = tensor
    elif ndim == 2:
        for ind, tensor in enumerate(tensor_list):
            all_tensor[ind, :tensor.shape[0], :tensor.shape[1]] = tensor
    elif ndim == 3:
        for ind, tensor in enumerate(tensor_list):
            all_tensor[ind, :tensor.shape[0], :tensor.shape[1], :tensor.shape[2]] = tensor
    else:
        raise Exception('Sorry, unimplemented.')
    return all_tensor
587557
def mean_with_attention(x, name, num_heads=4):
  """"""
  with tf.variable_scope(name):
    shape = shape_list(x)
    m = tf.reduce_mean(x, [1, 2])
    a = layers().Dense(num_heads, name="mean_attn")(x)
    s = tf.reshape(a, [shape[0], -1, num_heads])
    s = tf.nn.softmax(s, axis=1)
    s = tf.reshape(s, shape[:-1] + [1, num_heads])
    am = tf.reduce_mean(tf.expand_dims(x, axis=-1) * s, [1, 2])
    l = tf.concat([am, tf.expand_dims(m, axis=-1)], axis=-1)
    return layers().Dense(2 * shape[-1], name="mean_attn_final")(
        tf.reshape(l, [shape[0], (num_heads+1) * shape[-1]]))