How to take the first i-slices<traced> of a variable in a scan loop

[ziky168]

I passed a jnp array containing only one integer into a scan loop, in order to calibrate that this is the first few loops, just like for_i. But I had to write it in the form of scan.
For example, logprobs[i, :] the i is Traced<ShapedArray(float64[1]):JaxprTrace(level=1/0)>.
Then IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax.
Then, I tried jax.lax.switch,exactly:
def select_logprobs(self,i,logprobs): i = jnp.array(i).item() cases = [lambda: logprobs[j, :] for j in range(self.L)] return jax.lax.switch(i, cases)
For self.select_logprobs(i,logprobs) ,Error raised: *** TypeError: Branch index must be scalar, got Traced<ShapedArray(int32[1]):JaxprTrace(level=1/0)> of shape (1,).
For jax.jit(self.select_logprobs(i,logprobs),static_argnames=0) ,Error raised: *** TypeError: Branch index must be scalar, got Traced<ShapedArray(int32[1]):JaxprTrace(level=1/0)> of shape (1,).
How can I get a slice when the slice_size is a Traced<ShapedArray(float64[1]):JaxprTrace(level=1/0)>?
Thank you all for your selfless help!!!

[jakevdp]

Hi - it's hard to figure out what code you're executing from the description you gave. Could you edit your question and add a minimal reproducible example that shows the issue you're having, so that we can better answer your question? Your question would also be clearer with properly formatted code; refer to Creating and highlighting code blocks for information on that. Thanks!

[ziky168]

Sure! Thank you!!!

There are an minimal example:

import jax
from jax.config import config
config.update("jax_enable_x64", True)
import flax.linen as nn
import jax.numpy as jnp
import flax
from typing import Union
from functools import partial

def cplx_variance_scaling(rng, shape, dtype):
    rng1, rng2 = jax.random.split(rng)
    unif = jax.nn.initializers.uniform(scale=1.)
    elems = 1
    for k in shape[:-2]:
        elems *= k
    w = jax.numpy.sqrt((shape[-1] + shape[-2]) * elems)
    return (1. / w) * unif(rng1, shape, dtype=jnp.float64) * jax.numpy.exp(1.j * 3.141593 * unif(rng2, shape, dtype=jnp.float64))

def init_fn_args(dtype=None, **kwargs):
    init_args = {}
    if dtype is not None:
        init_args["dtype"] = dtype
    if flax.__version__ >= "0.4.0":
        init_args["param_dtype"] = dtype
        for k in kwargs.keys():
            if k[-4:] == "init":
                init_args[k] = kwargs[k]
    else:
        for k in kwargs.keys():
            if k[-4:] == "init":
                init_args[k] = partial(kwargs[k], dtype=dtype)

    return init_args

class RNN1DGeneral(nn.Module):
    L: int = 10
    depth: int = 1
    inputDim: int = 2
    hiddenSize: int = 10
    actFun: callable = nn.elu
    initScale: float = 1.0
    logProbFactor: float = 1
    realValuedOutput: bool = False
    realValuedParams: bool = True
    cell: Union[str, list] = "NADE"
    initDist:str = None
    initialized:bool = True
    MConstrains: jnp.ndarray = jnp.array([2, 2, 21, 21]) 

    def setup(self):
        if isinstance(self.cell, str) and self.cell != "RNN":
            ValueError("Complex parameters for LSTM/GRU not yet implemented.")

        if self.realValuedParams:
            self.dtype = jnp.float64
            self.initFunction = jax.nn.initializers.normal(stddev=1)
        else:
            self.dtype = jnp.complex128
            self.initFunction = partial(cplx_variance_scaling, scale=self.initScale)

        if isinstance(self.cell, str):
            self.zero_carry = jnp.zeros((self.depth, 1, self.hiddenSize), dtype=self.dtype)
            if self.cell == "NADE":
                self.cells = [NADECell() for _ in range(self.depth)]
                self.zero_carry = jnp.zeros((self.depth, 1), dtype=jnp.int32)
            else:
                ValueError("Cell name not recognized.")
        else:
            self.cells = self.cell[0]
            self.zero_carry = self.cell[1]

        self.rnnCell = RNNCellStack(self.cells, actFun=self.actFun)
        
    def __call__(self, x):
        _, probs = self.rnn_cell((self.zero_carry, jnp.zeros(self.inputDim)[None,:],self.MConstrains), (jax.nn.one_hot(x, self.inputDim)))
        
        return jnp.sum(probs)

    @partial(nn.transforms.scan,
              variable_broadcast='params',
              split_rngs={'params': False})
    def rnn_cell(self, carry, x):
        # x1, MConstrain = x
        if self.cell =="NADE":
            newCarry, logProb = self.rnnCell(carry[0], carry[1])
        else:
            newCarry, out = self.rnnCell(carry[0], carry[1])
            logProb = self.log_coeffs_to_log_probs(self.outputDense(out),carry[2])
        logProb = jnp.sum(logProb * x, axis=-1)
        return (newCarry, x), jnp.nan_to_num(logProb, nan=-35)
    

class RNNCellStack(nn.Module):
    """
    Implementation of a stack of RNN-cells which is scanned over an input sequence.
    This is achieved by stacking multiple 'vanilla' RNN-cells to obtain a deep RNN.
    
    Arguments:
    * ``hiddenSize``: size of the hidden state vector
    * ``actFun``: non-linear activation function
    * ``initScale``: factor by which the initial parameters are scaled
    
    Returns:
    New set of hidden states (one for each layer), as well as the last hidden state, that serves as input to the output layer
    """

    cells: list
    dtype: type = jnp.float64
    actFun: callable = nn.elu
    initFun: callable = jax.nn.initializers.variance_scaling(scale=0.1, mode="fan_avg", distribution="normal")
    # initFun: callable = jax.nn.initializers.normal(stddev=0.01)

    @ nn.compact
    def __call__(self, carry, newR):
        newCarry = jnp.zeros_like(carry)
       
        newR = nn.Dense(features=carry.shape[-1], use_bias=False,
                        **init_fn_args(kernel_init=self.initFun, dtype=self.dtype), 
                        name="data_in_dense")(newR)
        newR = self.actFun(newR)

        for j, (c, cell) in enumerate(zip(carry, self.cells)):
            current_carry, newR = cell(c, newR)
            newCarry = newCarry.at[j].set(current_carry)
        return newCarry, newR
    

class NADECell(nn.Module):
    # actFun: callable = nn.elu
    actFun: callable = nn.sigmoid
    dtype: type = jnp.float32
    L: int = 10
    depth: int = 1
    inputDim: int = 2
    hiddenSize: int = 10

    def setup(self):
        self.W = self.param('W', jax.nn.initializers.normal(), (self.hiddenSize, self.L), self.dtype)
        self.c = self.param('c', jax.nn.initializers.zeros, (self.hiddenSize,), self.dtype)
        self.V = self.param('V', jax.nn.initializers.normal(), (self.L, self.inputDim, self.hiddenSize), self.dtype)
        self.b = self.param('b', jax.nn.initializers.zeros, (self.L, self.inputDim,), self.dtype)
    
    def select_logprobs(self,i,logprobs):
        cases = [lambda: logprobs[j, :] for j in range(self.L)]
        return jax.lax.switch(i, cases)
    
    @nn.compact
    def __call__(self, i, logprobs):
        probs = jnp.exp(jax.jit(self.select_logprobs(i,logprobs),static_argnames=0)) 
        states = jax.random.categorical(self.make_rng('params'), probs, axis=-1).reshape() 
        h_i = self.actFun(self.c + jnp.einsum("hi,bi->bh", self.W[:, :i], states))
        logits = jax.nn.log_softmax(self.b[i,:] + jnp.einsum("oh,bh->bo", self.V[i,:], h_i)) 
        ids = logits[:, i].astype(jnp.int64)
        log_prob = jnp.take_along_axis(logits, ids[:, None], axis=1).squeeze(1) 
        logprobs = jnp.concatenate([logprobs, log_prob[None, :]], axis=0)
        newCarry = (i + 1, logprobs)
        return newCarry, log_prob

The example code includes several functions, among which cplx_variance_scaling and init_fn_args serve as auxiliary functions and can be omitted. RNNCellStack is stacked up with NADEcell. When calling self.rnn_cell, an error occurs in the first line of the __call__ function in the NADECell class: I cannot access the first i slices of the variable logprobs (where i is Traced<ShapedArray(float64[1]):JaxprTrace(level=1/0)>, and the initially passed i is jnp.array([0])). I am completely clueless about how to handle variables within traced.

[jakevdp]

Can you edit your example to include the code that reproduces the error you're seeing (i.e. which class should I instantiate, with what data, and what methods should I call?)

[ziky168]

Sorry, my example seems quite complex, with three levels of nesting ( from low to high: NADECell < RNNCellStack < RNN1DGeneral ). The part that needs optimization is in the call function of the NADECell class. You can instantiate the RNN1DGeneral class with the following simple code, using the default parameters for setup;

input_data = jnp.array([0, 1], dtype=jnp.int32)
model = RNN1DGeneral()
rng = jax.random.PRNGKey(0)
params = model.init(rng, input_data)
output = model.apply(params, input_data)

when calling the call function of RNN1DGeneral, you will find an error occurring in the call function of the NADECell class, which is what I need to optimize.
The specific sentence that has an error is :

probs = jnp.exp(jax.jit(self.select_logprobs(i,logprobs),static_argnames=0)) 

Error raised:TypeError: Branch index must be scalar, got Traced<ShapedArray(int32[1]):JaxprTrace(level=1/0)> of shape (1,).
The function of this sentence is to achieve :

probs = jnp.exp(logprobs[i,:]) 

Error raised:IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax.
My jax version=0.3.15,jaxlib=0.3.14.

[jakevdp]

Thanks - the error I get when running your code is this:

TypeError: Branch index must be scalar, got Traced<ShapedArray(int32[1]):JaxprTrace(level=1/0)> of shape (1,).

This indicates that you're passing a length-1 array in a place where a scalar is expected. I fixed this by changing this:

        return jax.lax.switch(i, cases)

to this:

        return jax.lax.switch(i[0], cases)

Running again, I get this error:

TypeError: Expected a callable value, got Traced<ShapedArray(float64[1]):JaxprTrace(level=1/0)>

This comes because you're passing an array to JIT rather than a function. To fix this I changed this:

        probs = jnp.exp(jax.jit(self.select_logprobs(i,logprobs),static_argnames=0)) 

to this:

        probs = jnp.exp(jax.jit(self.select_logprobs,static_argnames=0)(i,logprobs)) 

Running it again, I get this error:

TypeError: 'int' object is not iterable

This indicates that you're passing an integer to a function that's expecting a sequence. To fix it I changed static_argnames=0 to static_argnums=0.

Running it again, I get this error:

IndexError: Array slice indices must have static start/stop/step to be used with NumPy indexing syntax. Found slice(None, Traced<ShapedArray(int32[1]):JaxprTrace(level=1/0)>, None). To index a statically sized array at a dynamic position, try lax.dynamic_slice/dynamic_update_slice (JAX does not support dynamically sized arrays within JIT compiled functions).

This comes from this expression: self.W[:, :i]. Here i is a traced value, so the result of self.W[:, :i] would be a dynamically shaped array, which is not supported in JAX (see JAX sharp bits: dynamic shapes).

Unfortunately there's no easy fix here: you cannot use dynamically shaped arrays in a JIT-compiled JAX model, so you'll have to rewrite your model in a way that uses static-shaped arrays. I'm not sure what to recommend, because I don't entirely understand the intent of your code: as written, states is a scalar, and it's not clear to me why you'd pass a scalar to einsum.

But generally speaking in situations like this, you can often use a padded representation in place of things like self.W[:, :i], potentially zeroing out relevant pieces of the subsequent computations to compute the same quantities without the use of dynamically-sized arrays.

Does that make sense?

[ziky168]

Thank you very much for your kind efforts. Perhaps I need to make some adjustments to the model structure.