Slice JAX ShapedArray using indices

[snat1505027]

Hi, I have a JAX ShapedArray, X, of shape (1, 1080, 768) -> (batch_size, seq_len, hidden_dim) and another sorted ShapedArray, IDX, of size (1, 270) where the values in the IDX represent the indices of seq_len in X. I want to create a ragged tensor, Y, where the values in the X will be replaced by MASK if that index appears in IDX.

For example, X = [[[ 9, 10, 3],
[ 7, 10, 0]]]
IDX = [[1]]

                  then Y = [[[ 9, 10,  3],[MASK]]]

Is there any efficient way to do it?

[jakevdp]

Can you say more about why you're working with ShapedArray objects? They are essentially just abstract descriptors that contain a shape and a dtype. They don't contain any data, so there's not really any way to "slice" them, or to reason about the values in them. Perhaps you could edit your question with a minimal executible example of the kind of thing you want to do?

[snat1505027]

import dataclasses

import jax
import jax.numpy as jnp
import numpy as np
import flax.linen as nn
from jax import lax
from typing import Any, Dict, Union, Optional
import math
from jax import jit
from copy import deepcopy
import clu.parameter_overview
from dataclasses import dataclass
from functools import partial
import tensorflow as tf
from jax import random

MASK = 3

batch_size = 1
vis_seq_len = 1080
hidden_dim = 768
mask_rate = 0.25

def uniform_random_select(n, num_samples, sort_idx=True):
    """
    Randomly choose "num_samples" from N
    :param n:
    :param num_samples:
    :param sort_idx: Whether to sort the resulting index
    :return:
    """
    if isinstance(num_samples, int) and isinstance(n, int):
        assert num_samples <= n
    logits = tf.random.uniform([n])
    idx = tf.argsort(logits)[:num_samples]
    if sort_idx:
        idx = tf.sort(idx)
    idx = tf.cast(idx, dtype=tf.int32)
    return idx

key = random.PRNGKey(758493)
image_enc = jax.random.uniform(key, [batch_size, vis_seq_len, hidden_dim])
num_image_spans = vis_seq_len
num_image_patches_trg = int(num_image_spans * mask_rate)
image_patch_trg_idx = jnp.asarray(uniform_random_select(n=num_image_spans, num_samples=num_image_patches_trg, sort_idx=True)[None])

def produce_tensor_with_mask_token(image_enc, image_patch_trg_idx, num_image_spans):
    mask_one_hot = tf.reduce_any(tf.equal(tf.range(num_image_spans)[:, None], image_patch_trg_idx), 1)[None]
    '''returns a ragged jax array 'new_image_inc' of shape (1, 1080, variable size) where 
	if mask_one_hot[x][y] = False
        	the new_image_inc[x][y] value will be same as the image_enc[x][y] value 
        else
        	the new_image_inc[x][y] value will be MASK token'''

Sorry again, does this work?

[jakevdp]

I see, thanks. JAX doesn't support ragged arrays (every row must be the same size), so one possibility would be to fill the full row with the MASK token. Would that work for your usecase?

[jakevdp]

I feel bad this took so long – I keyed in on the wrong thing on your original question (ShapedArray) where I should have keyed in on "ragged", which JAX doesn't support.

[snat1505027]

Thanks for the reply. I had another question. If I have a DeviceArray, X, of shape (945, 4, 768) -> (row, column, hidden_dim) and another DeviceArray, IDX, of shape (945, 1) where each entry in IDX is a random integer value between 0 to 3 indicating one of the column in each row of X. I want to return a new DeviceArray where for row p, the hidden_dim in the column mentioned in IDX[p] will be replaced by tf.zeros([768]).

import jax

rkey = jax.random.PRNGKey(758493)
seq_mask = jax.random.randint(rkey, [945, 1], minval=0, maxval=4)
seq = jax.random.uniform(rkey, [945, 4, 768])

def produce_masked_seq(seq, seq_mask):
      ''' returns new_seq with shape (945, 4, 768) where in each row of new_seq there will be one column with tf.zeros([768]) and that column will be given in seq_mask '''

Is there any in-line code for that? Please let me know if you feel the code/description is ambiguous. Thanks!

[jakevdp]

If I understand correctly, I think you're looking for something like out = jax.vmap(lambda x, i: x.at[i].set(0))(seq, seq_mask)