Zhixuan Lin

Reward clipping and value transformation

1

Hello, Thanks for this great work! I noticed that you choose to clip the reward to [-1, 1] for Atari. I'm wondering what's the purpose of applying value transformation (i.e....

First, thanks for this awesome codebase, it helped me a lot :) I have three questions 1. Based on #147 and the [white paper](https://arxiv.org/abs/1812.06110), the results in the [baseline folder](https://github.com/google/dopamine/tree/master/baselines)...

`attend_dtype` not used

1

[Here](https://github.com/google/maxtext/blob/5353a957594bd6cf316747cd5a327c163caca74f/MaxText/layers/embeddings.py#L102-L103) it seems that the hard-coded `bfloat16` is used instead of `attend_dtype`. Also `query` is not cast. I guess the correct behavior should be casting both `query` and `self.embedding` to...

simple pallas kernel hangs when input size exceeds some threshold

2

### Description The following simple pallas kernel that copies an array hangs indefinitely: ```python import jax import jax.numpy as jnp from jax.experimental import pallas as pl def copy_kernel( src, dst...